diff --git a/.gitattributes b/.gitattributes index a6544df53e0665ce298e2f292dcae68ae81f28d4..cbdeacc41b28b4288780e989f12f46fe73e769a6 100644 --- a/.gitattributes +++ b/.gitattributes @@ -39,3 +39,53 @@ models/ppo1-aim500/0/default/ocdbt.process_0/d/0e705ab1649291dcd4833d7c54c84c68 models/ppo1-aim500/0/default/ocdbt.process_0/d/1842d714394728bf2cd8a101ca51be79 filter=lfs diff=lfs merge=lfs -text models/ppo1-aim500/0/default/ocdbt.process_0/d/76b87323ea6fece1f3000f3671f7a65c filter=lfs diff=lfs merge=lfs -text models/ppo1-aim500/0/default/ocdbt.process_0/d/e127cf6e78fa29c319f572ccfd0c1d80 filter=lfs diff=lfs merge=lfs -text +models/ppo1/0/default/d/648c91d28d47c7fdf4ef44a1fdd03ce7 filter=lfs diff=lfs merge=lfs -text +models/ppo1/0/default/ocdbt.process_0/d/af265c2609ca8829fd3e63f04fe41305 filter=lfs diff=lfs merge=lfs -text +models/ppo1/0/default/ocdbt.process_0/d/eb91c04a3dabe1aac4233a3752f8ae83 filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-fixed-reward-estimate/0/default/d/15a85f646bafb55f51a18966a8a4eba3 filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/1691c69e471866bf5d174ba00029a7d2 filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/f2e7e5c758eb4773c9cee341213d7c29 filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-mi-only/0/default/d/317d4d2f6c3a4abcade3f68dc576f6a9 filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/6f764aa8da8d20f9790c34036f8ebf6c filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/719b2023adb97e8e004b975110d4c980 filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-novelty-only/0/default/d/8a54313f46686b4f8386e4a65f7f301b filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/271948bb29e4c6d53499572b3503318b filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/8e59f3ea1b09fe58e0cfcb486371a084 filter=lfs diff=lfs merge=lfs -text +models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/d8205362561c89b0d3e19fa06f18d00e filter=lfs diff=lfs merge=lfs -text +models/ppo1-dads500/0/default/d/d311620e8f63bc8eadb425e79cf903d1 filter=lfs diff=lfs merge=lfs -text +models/ppo1-dads500/0/default/ocdbt.process_0/d/5052985b5acf0be41884125bcbb1d120 filter=lfs diff=lfs merge=lfs -text +models/ppo1-dads500/0/default/ocdbt.process_0/d/ee22700a49c9fa9af6bd1b82129863a9 filter=lfs diff=lfs merge=lfs -text +models/ppo1-dads500/0/default/ocdbt.process_0/d/f289a04c53c6079360b246d05472938f filter=lfs diff=lfs merge=lfs -text +models/ppo1-diayn500/0/default/d/a1433643c777305b17c52c5376bc0c57 filter=lfs diff=lfs merge=lfs -text +models/ppo1-diayn500/0/default/ocdbt.process_0/d/265ace8b1373aa11596e3e75feedf482 filter=lfs diff=lfs merge=lfs -text +models/ppo1-diayn500/0/default/ocdbt.process_0/d/8ac165a32779d4aed6fc4892996b8dc3 filter=lfs diff=lfs merge=lfs -text +models/ppo1-diayn500/0/default/ocdbt.process_0/d/bd998024e8a75a5d0f62ffabf3051b3c filter=lfs diff=lfs merge=lfs -text +models/ppo1-icm500/0/default/d/5dc7664f7ddd729e30fe48716300c09a filter=lfs diff=lfs merge=lfs -text +models/ppo1-icm500/0/default/ocdbt.process_0/d/5e2e6028207e4d2739be514a4100684c filter=lfs diff=lfs merge=lfs -text +models/ppo1-icm500/0/default/ocdbt.process_0/d/de1203851f6380ed6e4899697f68e66a filter=lfs diff=lfs merge=lfs -text +models/ppo1-metra500/0/default/d/fd19d4383cd42974467aeb5bf1554a4c filter=lfs diff=lfs merge=lfs -text +models/ppo1-metra500/0/default/ocdbt.process_0/d/1d9d018ec71450f96ca3628841948c3d filter=lfs diff=lfs merge=lfs -text +models/ppo1-metra500/0/default/ocdbt.process_0/d/ec10c9c8cf0eb47449729e032e204e75 filter=lfs diff=lfs merge=lfs -text +models/ppo1-random-walk500/0/default/d/036e4c234628f458d5f7a9847cace240 filter=lfs diff=lfs merge=lfs -text +models/ppo1-random-walk500/0/default/ocdbt.process_0/d/2077fa6cda20d4df8af93326760d3a2b filter=lfs diff=lfs merge=lfs -text +models/ppo1-random-walk500/0/default/ocdbt.process_0/d/d9ed0ffbb5d2cf78cf8a567878327ad6 filter=lfs diff=lfs merge=lfs -text +models/ppo1-recovery500/0/default/d/d9223f38b238fa1d02f4eb96914f5fac filter=lfs diff=lfs merge=lfs -text +models/ppo1-recovery500/0/default/ocdbt.process_0/d/14825823161b2bacd0a1449ca806f4a7 filter=lfs diff=lfs merge=lfs -text +models/ppo1-recovery500/0/default/ocdbt.process_0/d/b49d2426d299be18a6217211a14c78ad filter=lfs diff=lfs merge=lfs -text +models/ppo1-recovery500/0/default/ocdbt.process_0/d/eb0be67ba49e905ab08a513dad7d0573 filter=lfs diff=lfs merge=lfs -text +models/ppo1-uniform-random500/0/default/d/fff057e04911c88885541c50ff03e80f filter=lfs diff=lfs merge=lfs -text +models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/10b5e13311f1fe474a649fff5059396e filter=lfs diff=lfs merge=lfs -text +models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/f6f509b14bfe5f51f56c22ccb3910b58 filter=lfs diff=lfs merge=lfs -text +models/ppo100/0/default/d/0300e5ac4b49bffe5719d9cfb4762d1a filter=lfs diff=lfs merge=lfs -text +models/ppo100/0/default/ocdbt.process_0/d/3c2bd48694a308b3376ff141afe1a568 filter=lfs diff=lfs merge=lfs -text +models/ppo100/0/default/ocdbt.process_0/d/c08f0db736a2a786d6c2b6a5f8be193d filter=lfs diff=lfs merge=lfs -text +models/recovery1/0/default/d/d892c60f4654593696558e2196fb3f0f filter=lfs diff=lfs merge=lfs -text +models/recovery1/0/default/ocdbt.process_0/d/30c5f0e29e4f4e5045850e4ae8ac8c14 filter=lfs diff=lfs merge=lfs -text +models/recovery1/0/default/ocdbt.process_0/d/7c0764d7c8f8db23a57c6d278af3966f filter=lfs diff=lfs merge=lfs -text +models/recovery1/0/default/ocdbt.process_0/d/ae2fbcfef5c278ec27aa5ca43aa085ad filter=lfs diff=lfs merge=lfs -text +models/recovery1-transferred-aim500/0/default/d/fc2640fe3dc160b15796bb4e93702707 filter=lfs diff=lfs merge=lfs -text +models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/262859f5412e92128ec9d77f68d5849e filter=lfs diff=lfs merge=lfs -text +models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/362947cf928ae683427fe0c00bfe6855 filter=lfs diff=lfs merge=lfs -text +models/recovery100/0/default/d/83878ea0eec4a9c18daeb44b066196fb filter=lfs diff=lfs merge=lfs -text +models/recovery100/0/default/ocdbt.process_0/d/77cebf3655ec723496319cdb1c72edbc filter=lfs diff=lfs merge=lfs -text +models/recovery100/0/default/ocdbt.process_0/d/d1816013dc8ccef6782787251339f26d filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index 555a342f2318ddb666ffd973ee6e13497df30ece..73be1472f88c552fc8ca11439f894fca4a583554 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ behavior; they are not an aggregate benchmark or a ranking of models. ## Models at a glance -**1 checkpoint · 256 × 256 RGB · approximately 580 MB each · 0.58 GB total** +**16 checkpoints · 256 × 256 RGB · approximately 580 MB each · 9.28 GB total** - **Inputs:** RGB observations and continuous 6D end-effector delta actions. - **Outputs:** predicted RGB frames at 256 × 256 resolution. @@ -52,6 +52,21 @@ collection iterations, with a 2:1 PPO-to-AIM sampling mixture. | Model | Training step | Download | | --- | ---: | ---: | | [PPO1 + AIM500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500) | 36,000 | 580.1 MB | +| [PPO1](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1) | 26,000 | 580.0 MB | +| [PPO1 + AIM500 (fixed reward estimate)](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500-fixed-reward-estimate) | 42,000 | 580.0 MB | +| [PPO1 + AIM500 (MI only)](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500-mi-only) | 39,000 | 580.0 MB | +| [PPO1 + AIM500 (novelty only)](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-aim500-novelty-only) | 42,000 | 580.1 MB | +| [PPO1 + DADS500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-dads500) | 29,000 | 580.0 MB | +| [PPO1 + DIAYN500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-diayn500) | 34,000 | 580.1 MB | +| [PPO1 + ICM500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-icm500) | 29,000 | 579.9 MB | +| [PPO1 + METRA500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-metra500) | 26,000 | 580.0 MB | +| [PPO1 + random walk500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-random-walk500) | 28,000 | 580.0 MB | +| [PPO1 + Recovery500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-recovery500) | 27,000 | 580.1 MB | +| [PPO1 + uniform random500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo1-uniform-random500) | 33,000 | 580.0 MB | +| [PPO100](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/ppo100) | 44,000 | 580.1 MB | +| [Recovery1](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/recovery1) | 37,000 | 580.0 MB | +| [Recovery1 + transferred AIM500](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/recovery1-transferred-aim500) | 42,000 | 580.1 MB | +| [Recovery100](https://huggingface.co/smbml/future-sight-pusht-models/tree/main/models/recovery100) | 40,000 | 580.1 MB | The **fixed reward estimate** ablation uses AIM data without reward re-estimation. Novelty-only and MI-only models ablate the reward components. @@ -126,7 +141,7 @@ a license to the separately trained Future Sight weights. ## Validation and limitations -All 16 source checkpoints passed strict restoration and finite-output checks on two selected evaluation cases each. They also ran in the local interactive model switcher. Published model files are listed above. +All 16 source checkpoints passed strict restoration and finite-output checks on two selected evaluation cases each. They also ran in the local interactive model switcher. Published model files are listed above. The PPO1 + AIM500 bundle additionally passed an authenticated Hugging Face upload/download round trip, file-hash verification, offline cache reuse, and exact raw-frame agreement with its local source on BC and PushAway when independently restored states use the same compiled inference graph. Separately compiled GPU runs were not pixel-identical; this check does not establish bitwise reproducibility across runs. Other models have local inference checks and remote file verification, not separate HF round-trip inference runs. The local checks used JAX/JAXLIB 0.10.2, Flax 0.12.7 and Orbax 0.12.1 on an NVIDIA RTX PRO 6000 Blackwell GPU. They do not establish compatibility with every diff --git a/models.json b/models.json index 0eb3b9597e4e6e4af702c60d9f1e710f71ea9a2d..cdc73e58cd6be11685f08f86189ae859e471a8b6 100644 --- a/models.json +++ b/models.json @@ -83,6 +83,1181 @@ "sha256": "57078aa520fa28dba100d56593b1a42ab6c58b0c3b631d3bbdae929282d59cf5" } ] + }, + { + "id": "ppo1", + "label": "PPO1", + "path": "models/ppo1", + "state_root_sha256": "5d1cc061fd9ab6a368c3b4670b1efbfd3b183c72bf828828fc147e56a2024539", + "training_step": 26000, + "weights": "ema", + "size_bytes": 579983383, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "41c46df8a8b898dec6c915418378b279e71657022ec7df7147c079a57973fc0f" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/648c91d28d47c7fdf4ef44a1fdd03ce7", + "size_bytes": 112804, + "sha256": "c1b4b3c8e99ee6a4fcc9b2c359fb194399ca2f63c0df74cb6b698a250ce8223b" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "031daa9457fe29469b3ddf061121e46061f9110bea47f025293eb2229bc7b4c8" + }, + { + "path": "0/default/ocdbt.process_0/d/0714fab6dc2ad102a2e15c141a03257b", + "size_bytes": 186, + "sha256": "78f1963429d958c839b0072dea545e0c9e8bc3bbf61e0627b723f9f497d79569" + }, + { + "path": "0/default/ocdbt.process_0/d/a10a646ba0743686f0f1c30ee71174d4", + "size_bytes": 2060, + "sha256": "fdb904682fae02b2411046985abf38e744b0d1030392584875bb6b44c4043019" + }, + { + "path": "0/default/ocdbt.process_0/d/af265c2609ca8829fd3e63f04fe41305", + "size_bytes": 216375296, + "sha256": "1a889568c15d391b40c767f95a92ee4cac3f6b6c34221cbe8324a3f2f517ed78" + }, + { + "path": "0/default/ocdbt.process_0/d/dc664153567cef23fe40a96f24e05ecc", + "size_bytes": 13502, + "sha256": "9b7ba50ee9a9da4ea1baeb44cbc13d4978dd072988e394bfa90960ec3673e8ea" + }, + { + "path": "0/default/ocdbt.process_0/d/eb91c04a3dabe1aac4233a3752f8ae83", + "size_bytes": 362860544, + "sha256": "26b075c7ef67b8ecd7bbb95c99f297f2303fc0de8ed2b3f847db012f23eeec87" + }, + { + "path": "0/default/ocdbt.process_0/d/eeffb558247f75dea678f8cfc8d41f63", + "size_bytes": 1998, + "sha256": "a03930149c9aa0ac2d496dec3781d54833f2ca218d1f8fdea69fade69c170877" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 353, + "sha256": "00b84a825528abdc838721c34a481c55534de7ff3dfb941c08bec9d0aa935fd7" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "866becfb8686dfa6895eb824f3a187b32b85f6c1234d2f78717b85abefb26987" + } + ] + }, + { + "id": "ppo1-aim500-fixed-reward-estimate", + "label": "PPO1 + AIM500 (fixed reward estimate)", + "path": "models/ppo1-aim500-fixed-reward-estimate", + "state_root_sha256": "24780d3293d74457c42332d4a1f63b2fb3cb18295f6e7796af644826e3d27cbb", + "training_step": 42000, + "weights": "ema", + "size_bytes": 579972718, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "4bf3dc3eb06b8d13ce991a01a52b76ed7ce529970df1d862f210f3d6ad75f2e8" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/15a85f646bafb55f51a18966a8a4eba3", + "size_bytes": 112946, + "sha256": "6b16cd6917850e5d2fb161c7b94ea3fdd4b542bd928f05e96b1884a575c04a4c" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "e2ee31b75960199daea2dbda6d1e6b7519860294a21c29f5b80999d0ba3119b9" + }, + { + "path": "0/default/ocdbt.process_0/d/1691c69e471866bf5d174ba00029a7d2", + "size_bytes": 541900800, + "sha256": "5e9efd37754a3b542b24415f98c2355e3cada3c6da515eaffb3e31514be9ea90" + }, + { + "path": "0/default/ocdbt.process_0/d/6cdfe6d8112774ebea68968a8b8b2450", + "size_bytes": 228, + "sha256": "5fc5c68027d1a329fd52df5ed3d021c8adfb02abd7da3392c521d8c91669047b" + }, + { + "path": "0/default/ocdbt.process_0/d/7931d26a51fcf24844ca857b5b1775cb", + "size_bytes": 13489, + "sha256": "15658532440577e8c426c63e23d57d849a62adf0de037e8570631625a58f646d" + }, + { + "path": "0/default/ocdbt.process_0/d/a5c70bae84187efd011e1c2d1ee4be1e", + "size_bytes": 1456, + "sha256": "97373f363fe79ac6d957109ceea688e93601b7e83d06a12ec607a42dbc094275" + }, + { + "path": "0/default/ocdbt.process_0/d/f2e7e5c758eb4773c9cee341213d7c29", + "size_bytes": 37326848, + "sha256": "f3c66b0ca613fd2df0d7838e15033b8f2bf4e0e2dc05143d467b730c202610ff" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 311, + "sha256": "fefcd617090b2e01a63cc37eaafa4c17754bacb116648a51b69a65cddeab28b4" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "9cd62d514eec8f8c63efbc443524ac5047bc27a1b31ca136f203a8f78673eb58" + } + ] + }, + { + "id": "ppo1-aim500-mi-only", + "label": "PPO1 + AIM500 (MI only)", + "path": "models/ppo1-aim500-mi-only", + "state_root_sha256": "97d58c357783468d43a5a6213d2eebb962c7f1d3aea9d24f67e883d272738155", + "training_step": 39000, + "weights": "ema", + "size_bytes": 579982203, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "1ae688c416b31c433fdc81c007c269f4f7870bbee0b475a1fcdc5061926b1a79" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/317d4d2f6c3a4abcade3f68dc576f6a9", + "size_bytes": 112873, + "sha256": "6f9717e3c8590a82d1bf8d6e82ec27c9fe5a19c7d3fff459216145b2435012a8" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "87fde8ef67778090591f584dcdaa249d812d9f0e7fe9e7dea825b464bda9a51e" + }, + { + "path": "0/default/ocdbt.process_0/d/6f764aa8da8d20f9790c34036f8ebf6c", + "size_bytes": 59101184, + "sha256": "798d4f8458ecc45374241563f68149da25d5e7ac96cf560b2194054a9643c6a9" + }, + { + "path": "0/default/ocdbt.process_0/d/719b2023adb97e8e004b975110d4c980", + "size_bytes": 520146944, + "sha256": "792175f73e1473d2c4a53fc291ff7c03855f6950d9d8b553c078407842bd0890" + }, + { + "path": "0/default/ocdbt.process_0/d/960c0e6c50a75b95799bb0d1f70cfe29", + "size_bytes": 2060, + "sha256": "fdb904682fae02b2411046985abf38e744b0d1030392584875bb6b44c4043019" + }, + { + "path": "0/default/ocdbt.process_0/d/b8ab8a019fdf3cdfb451965c47d6b600", + "size_bytes": 186, + "sha256": "78f1963429d958c839b0072dea545e0c9e8bc3bbf61e0627b723f9f497d79569" + }, + { + "path": "0/default/ocdbt.process_0/d/ecbe1ddeac541714c7d1c9627edaa576", + "size_bytes": 1998, + "sha256": "a03930149c9aa0ac2d496dec3781d54833f2ca218d1f8fdea69fade69c170877" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 318, + "sha256": "aa95f65bf372ab61518870a67946e540934d64cf0db0fa9d8b07bc8a7deef903" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "2531e10b965b8aff25351628374c0990607766f8cd91fd196ecc49595f7f8897" + } + ] + }, + { + "id": "ppo1-aim500-novelty-only", + "label": "PPO1 + AIM500 (novelty only)", + "path": "models/ppo1-aim500-novelty-only", + "state_root_sha256": "a5f0cd0062bb7f90d774ca99d8c975366a022e93a3c99a9225f2d5351456e8b6", + "training_step": 42000, + "weights": "ema", + "size_bytes": 580114496, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "925b4e786ba2fe02ae3b210a03576051fd23787c2ddc12feb5c3b0b077947f8d" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/8a54313f46686b4f8386e4a65f7f301b", + "size_bytes": 112877, + "sha256": "196545e3a7629f25f6cc74f51f2f4e6acb54fef851cd315146b5dd1b0222873e" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "fd200dcbebbc91dd0a95be8216f60d100bd02962d5a51119504b46d7707200a4" + }, + { + "path": "0/default/ocdbt.process_0/d/271948bb29e4c6d53499572b3503318b", + "size_bytes": 514379776, + "sha256": "c3c853c1521113f08dd93ee62f610d912cb715b6b1094b7b4164611201b1d110" + }, + { + "path": "0/default/ocdbt.process_0/d/545a2abeb6a152cbaafcc1ad5aea98d2", + "size_bytes": 2096, + "sha256": "f27f5213b11e8447c6a86382042f9c96bdc59f3a907f6e066dd7f525e9deab74" + }, + { + "path": "0/default/ocdbt.process_0/d/583b331c85dcea84777c905344f3c9ce", + "size_bytes": 199, + "sha256": "77e6f800feb9f959ba9c8c693b4f6f183a9c090a693f4359f6381f405287634f" + }, + { + "path": "0/default/ocdbt.process_0/d/8e59f3ea1b09fe58e0cfcb486371a084", + "size_bytes": 61173760, + "sha256": "7654dd2b3c7f423957ebd95724aaa406bf1ca6c0b2b69a73d85028e99001f9e6" + }, + { + "path": "0/default/ocdbt.process_0/d/d8205362561c89b0d3e19fa06f18d00e", + "size_bytes": 3827385, + "sha256": "ff362576b4387b4f6911da1bbb939177af64ece59b9f9c6259c1cb5e9a4cf446" + }, + { + "path": "0/default/ocdbt.process_0/d/e1b1f6ccfb6fc53d6c9d307e38512ff2", + "size_bytes": 1401, + "sha256": "0d7638c33a0ad4331ed4e8d22466f5add97bdb294abe3ada60d3fcbcac2184ea" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 362, + "sha256": "5fab40fbc9ae29a535d61eb484b400e6d58e2c3748f5bbd26e394590248400ca" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "c46886ccb8aec2aa2017ec149ddaae69177bc46d742f69e86c8abb83576b396b" + } + ] + }, + { + "id": "ppo1-dads500", + "label": "PPO1 + DADS500", + "path": "models/ppo1-dads500", + "state_root_sha256": "288e3c30c44f576b6b86a821d75800fbb184ca5db86de77e3b78a8057c5f18cb", + "training_step": 29000, + "weights": "ema", + "size_bytes": 580039360, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "832c332b24ec4dd24d6e3b22a5b56e0f8a62d128431f73d2ca37b86f793de03a" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/d311620e8f63bc8eadb425e79cf903d1", + "size_bytes": 112879, + "sha256": "77044f7a39188f39af0e5ffc30d69633dc03809cde07524e7091ca6fe4a69470" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "0338d925b8cf9bcb15bea63b6d1441ab32fb976b6a5f82fd4ebf42ffdcf0f9de" + }, + { + "path": "0/default/ocdbt.process_0/d/3ea8182f4a879390c95bbc7a03ea0c38", + "size_bytes": 2060, + "sha256": "fdb904682fae02b2411046985abf38e744b0d1030392584875bb6b44c4043019" + }, + { + "path": "0/default/ocdbt.process_0/d/5052985b5acf0be41884125bcbb1d120", + "size_bytes": 359825408, + "sha256": "2b4e61a9f890ed1d9eaadc4937f9907fa9aa84a555a3bf5f6fe64ca7376578f0" + }, + { + "path": "0/default/ocdbt.process_0/d/c4a13bf5b23b4f3a8bba02e147f22930", + "size_bytes": 1763, + "sha256": "42f7625085774a95fb0671528afed5f390c076872264248a7e2b661cc1c91c44" + }, + { + "path": "0/default/ocdbt.process_0/d/dfd4b6d00ae3a8b91abf37c969f53f9f", + "size_bytes": 186, + "sha256": "78f1963429d958c839b0072dea545e0c9e8bc3bbf61e0627b723f9f497d79569" + }, + { + "path": "0/default/ocdbt.process_0/d/ee22700a49c9fa9af6bd1b82129863a9", + "size_bytes": 4382720, + "sha256": "e009ddcf61627bdfc7dc856c5e36f77d73e9822ecfdabcd1fde3a3da167b23a1" + }, + { + "path": "0/default/ocdbt.process_0/d/f289a04c53c6079360b246d05472938f", + "size_bytes": 215097344, + "sha256": "86df92ce262963b74fabee9fc20b4608722facedabcac49d3474937f68f5f6b3" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 360, + "sha256": "a403fd11013b044183d4b51ed0cba65c410b8e21e7cc9faa7e33ce3199ab8250" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "1b308328b98e88e182180d2a6987b52ccd1dfe62d2894b6f4cb3759dcd153179" + } + ] + }, + { + "id": "ppo1-diayn500", + "label": "PPO1 + DIAYN500", + "path": "models/ppo1-diayn500", + "state_root_sha256": "ff46161bc43e2b4474350cb98a07d71639839fe36e93fff76537c7517e9134d1", + "training_step": 34000, + "weights": "ema", + "size_bytes": 580072176, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "4ff77db841de572cfdc4a50508f3356ecbc06ce05bc1bd369754fd9fc6c80c42" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/a1433643c777305b17c52c5376bc0c57", + "size_bytes": 112878, + "sha256": "553b738b9a783d47f1f40444b7dda8b079ffcacc7547d44f8b642eea1f6541c6" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "ed9c84c8b02e061d28369abaeb7156a28853e1873f37a9b9495177e57b42ce34" + }, + { + "path": "0/default/ocdbt.process_0/d/265ace8b1373aa11596e3e75feedf482", + "size_bytes": 72097792, + "sha256": "da52e57c861351fdca66cfdadcee2e2a8cad29af5bfd5dcec887a497fb5a489a" + }, + { + "path": "0/default/ocdbt.process_0/d/3fc5aabae3972fda1850bdd8d2de913c", + "size_bytes": 186, + "sha256": "5b172a816dc68742ff2e673cf506cad549b3a2305f55c7f45a9bad56b90091f3" + }, + { + "path": "0/default/ocdbt.process_0/d/860e7ec4b2f6d773e351ef0565846806", + "size_bytes": 2057, + "sha256": "4834fc7fbc3880dca62b712b4e4c90d69669fcdc6297b5ac382bd52c41336364" + }, + { + "path": "0/default/ocdbt.process_0/d/8ac165a32779d4aed6fc4892996b8dc3", + "size_bytes": 17539072, + "sha256": "9e02525b32a92c81a51d5886ecf9ddec178eec06d704876cba398b31864ec50c" + }, + { + "path": "0/default/ocdbt.process_0/d/bd998024e8a75a5d0f62ffabf3051b3c", + "size_bytes": 489701376, + "sha256": "36e0443daad45d7d4d78ec806670ae5d38a5a5c2efe29056c47869e0ee38f21e" + }, + { + "path": "0/default/ocdbt.process_0/d/ee12ad97855496c5a52dce9b8e4e4e4f", + "size_bytes": 1814, + "sha256": "524754f24043fb0ae635d9568db6bed86df40f36f86bcdaab67a79150467c64b" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 361, + "sha256": "5eb124ea0da8046463aa298f9a12f47dcf38c5eaa0c03832cf20a15f9b556cb6" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "147030b408443c590a801b4b583cd62b647dbba536031168b16458535fd67df2" + } + ] + }, + { + "id": "ppo1-icm500", + "label": "PPO1 + ICM500", + "path": "models/ppo1-icm500", + "state_root_sha256": "9873b1fdafef2e8652ec2eacb7515e9a897b50341a812ad6da34cccda6038540", + "training_step": 29000, + "weights": "ema", + "size_bytes": 579949499, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "90218b1fb310302636b547498ae29f0137c18754349ddd7a542b3a3243bbab29" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/5dc7664f7ddd729e30fe48716300c09a", + "size_bytes": 112942, + "sha256": "209f629934f74ee01d125d3783cabfe1e99ff327de24d6d8c034ed32b3b2cc40" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "c1acc911edc6e18c3f766d9f43434b652406f2e52bdebd46420d51ab6f658c7f" + }, + { + "path": "0/default/ocdbt.process_0/d/18763adcffc7b993ea4faafcfc8e48f4", + "size_bytes": 1998, + "sha256": "a03930149c9aa0ac2d496dec3781d54833f2ca218d1f8fdea69fade69c170877" + }, + { + "path": "0/default/ocdbt.process_0/d/4e996b880df24475f131ce5b73da7f6c", + "size_bytes": 2060, + "sha256": "fdb904682fae02b2411046985abf38e744b0d1030392584875bb6b44c4043019" + }, + { + "path": "0/default/ocdbt.process_0/d/5e2e6028207e4d2739be514a4100684c", + "size_bytes": 526053376, + "sha256": "3d7b78ab14c96d38020560ba0f56638e9526b8e4a3b1fc46fa605b4ec897b130" + }, + { + "path": "0/default/ocdbt.process_0/d/7283cfb568f3403a47b5f10874196fbc", + "size_bytes": 186, + "sha256": "78f1963429d958c839b0072dea545e0c9e8bc3bbf61e0627b723f9f497d79569" + }, + { + "path": "0/default/ocdbt.process_0/d/de1203851f6380ed6e4899697f68e66a", + "size_bytes": 53161984, + "sha256": "966ca3d2af1f85898766c03581e59c3f8dcc47366007bc222d8ba9d6b6313196" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 313, + "sha256": "db49ecbc6f22df62eee2f392ba370c44e91b68b9ec25a62e3ee36f9e6f7964d5" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "42cb42a9d19d4d8fc8d28f0316a12cf6fb9542c5740b05ddd2490ff60729d021" + } + ] + }, + { + "id": "ppo1-metra500", + "label": "PPO1 + METRA500", + "path": "models/ppo1-metra500", + "state_root_sha256": "aa1741e5d8ae545e72a4ef185e2c8cb53ce7f0efd1b542e5ff76a052e5a5b33c", + "training_step": 26000, + "weights": "ema", + "size_bytes": 579961031, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "71aa5bafb8c113032fae3ed8688c3d893b045deb6e38a36569061b65f4639ab2" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/fd19d4383cd42974467aeb5bf1554a4c", + "size_bytes": 112860, + "sha256": "0a54f5f9c05ddb5b9e2b9e91f6f96e6f4bed5b503c8ddc8212ecfef0bf31ef83" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "5e8388fb24f8e087960d35db48bd3f86b9fa2b79dba232fceddc29afd21ea200" + }, + { + "path": "0/default/ocdbt.process_0/d/1d9d018ec71450f96ca3628841948c3d", + "size_bytes": 511373312, + "sha256": "ac4f76f6a56f58f938ec21f54e4ba3ea9f8b7499145ca3af39299b3c935c8164" + }, + { + "path": "0/default/ocdbt.process_0/d/79b3d14665e9343193c53e8fb237c828", + "size_bytes": 186, + "sha256": "78f1963429d958c839b0072dea545e0c9e8bc3bbf61e0627b723f9f497d79569" + }, + { + "path": "0/default/ocdbt.process_0/d/c4cd899b11e2b55f75cad9e4477e26a5", + "size_bytes": 1327, + "sha256": "0535f41cb5f1f8bcdf3f8c526f3e9ad629b8b2beff7a78594a55b348d1fc4def" + }, + { + "path": "0/default/ocdbt.process_0/d/ec10c9c8cf0eb47449729e032e204e75", + "size_bytes": 67854336, + "sha256": "a4e510ce0ac6d5abb66312b57493e5c9470e33b8e10d4d7717c9a2c597cb4889" + }, + { + "path": "0/default/ocdbt.process_0/d/f03e4e2e123e6f9d57c6f3ab1559f880", + "size_bytes": 2060, + "sha256": "fdb904682fae02b2411046985abf38e744b0d1030392584875bb6b44c4043019" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 310, + "sha256": "aae0a14dbde6b4539ba25d45c21cc62896f03f646a9075759ac856d57253aedd" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "b789a2b7f341ff71545ac4d67fc72a350b1250ac3661d146dcb5530ded5e6a15" + } + ] + }, + { + "id": "ppo1-random-walk500", + "label": "PPO1 + random walk500", + "path": "models/ppo1-random-walk500", + "state_root_sha256": "a25ff767f446db73456661acc6139fef6a30502b524f5cfb02924bef3272fc68", + "training_step": 28000, + "weights": "ema", + "size_bytes": 580002639, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "3010e0cea691da51e56134d819765600836025d3c89b052d8a69a5d7fe4c1514" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/036e4c234628f458d5f7a9847cace240", + "size_bytes": 112833, + "sha256": "04a1982ee5c5175ba46d0b07133813906e48565d2a91341ba1d47a2220e1d32c" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "f8d4ad4d659074d9f13fe15c470d8e177a7de468285c8c26c9414ad14bec04f1" + }, + { + "path": "0/default/ocdbt.process_0/d/13cdb3e1f6a7eb1b28d75f2a4818f613", + "size_bytes": 1998, + "sha256": "a03930149c9aa0ac2d496dec3781d54833f2ca218d1f8fdea69fade69c170877" + }, + { + "path": "0/default/ocdbt.process_0/d/2077fa6cda20d4df8af93326760d3a2b", + "size_bytes": 164593664, + "sha256": "4ba688a02e75852494fe5e22c716aaf94df055bce966c9f8dbf33c9d747e3d34" + }, + { + "path": "0/default/ocdbt.process_0/d/3cc37a3501a39d61c841c613b1e8e8d9", + "size_bytes": 2060, + "sha256": "fdb904682fae02b2411046985abf38e744b0d1030392584875bb6b44c4043019" + }, + { + "path": "0/default/ocdbt.process_0/d/af9775c73bcd886b3f09f7387af537eb", + "size_bytes": 186, + "sha256": "78f1963429d958c839b0072dea545e0c9e8bc3bbf61e0627b723f9f497d79569" + }, + { + "path": "0/default/ocdbt.process_0/d/d9ed0ffbb5d2cf78cf8a567878327ad6", + "size_bytes": 414674944, + "sha256": "ad6a9ae88893b28e6e0317d98407123fc05c7ef6a9e71ee7f491420ff8b29089" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 314, + "sha256": "5c243d8096d4ce9a1bc6779f100094306db025f66c8052ec3a6c98c3ecb2389a" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "dfeadd5679f85764e8d4e75600fa4c45bfedad549ab65bfb937b98cd7e5c0ca0" + } + ] + }, + { + "id": "ppo1-recovery500", + "label": "PPO1 + Recovery500", + "path": "models/ppo1-recovery500", + "state_root_sha256": "c0d80ece428a6e1f02438cc9ce1b88dc427e68e01636bd663350890f7e013ba7", + "training_step": 27000, + "weights": "ema", + "size_bytes": 580079947, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "a5242b60f9aa402a7f7b44fe55429fd46d9c80b608dfb7feaae39478f713794c" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/d9223f38b238fa1d02f4eb96914f5fac", + "size_bytes": 112897, + "sha256": "5d22e46d003f1d81909623bcbe6c2852093ff8240ee2a2b9d2a8791ecc74a242" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "386c5cabcda3332b135f486fa658ffc296bc1dfa224c80679e2d55497c371ccc" + }, + { + "path": "0/default/ocdbt.process_0/d/0bb4e45cd475cc4e93c21893284ec6c7", + "size_bytes": 186, + "sha256": "5b172a816dc68742ff2e673cf506cad549b3a2305f55c7f45a9bad56b90091f3" + }, + { + "path": "0/default/ocdbt.process_0/d/14825823161b2bacd0a1449ca806f4a7", + "size_bytes": 2236142, + "sha256": "a8ab8d8146a9820e5b2ef7e1791437d241aaffe4d68d85c398080f0ff4bf3d3d" + }, + { + "path": "0/default/ocdbt.process_0/d/b49d2426d299be18a6217211a14c78ad", + "size_bytes": 62070784, + "sha256": "7100554192b36cce3fcd80399dbe3af0265a59610943674364aa1c3863169b0d" + }, + { + "path": "0/default/ocdbt.process_0/d/eb0be67ba49e905ab08a513dad7d0573", + "size_bytes": 515039232, + "sha256": "3773301888aee857dc6c2152258cdc0c47b582f8e9d49c29430dd28540d31ba3" + }, + { + "path": "0/default/ocdbt.process_0/d/f1f019f57442b0d4229381f27b4c5f6f", + "size_bytes": 2057, + "sha256": "4834fc7fbc3880dca62b712b4e4c90d69669fcdc6297b5ac382bd52c41336364" + }, + { + "path": "0/default/ocdbt.process_0/d/fa2d110d5b9976dc317c2e9eceab8f09", + "size_bytes": 1645, + "sha256": "a579dc7022b802da1e039f05f3d859a120bf1b8a757de83d070a0c6af9d22c45" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 364, + "sha256": "8b078b04583b2e1c1b589aa78debfb8974300b736fb284355299aeca2cfede56" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "e89f9de667345f14d6002c1e97dece533048affa8b8f528f75e2e8536e48da6f" + } + ] + }, + { + "id": "ppo1-uniform-random500", + "label": "PPO1 + uniform random500", + "path": "models/ppo1-uniform-random500", + "state_root_sha256": "9e12c8634a545a517bd0fc80536133e0c063494cc5c62ebedc2e1099cb11eaa4", + "training_step": 33000, + "weights": "ema", + "size_bytes": 580001470, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "43c124f4b3d8f281ef0bc4797158ec5ffa6426df6c1db64af824473febecf875" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/fff057e04911c88885541c50ff03e80f", + "size_bytes": 112872, + "sha256": "467fd9004504edf6b442cba71db45e5cd8afa4b7bdcdb4138ef901e06af9ccab" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "d1c69d6fc23fe1f5c3fa80280e1040cb6917db27baa258325b4e3a7c3afe8e8d" + }, + { + "path": "0/default/ocdbt.process_0/d/10b5e13311f1fe474a649fff5059396e", + "size_bytes": 192405504, + "sha256": "42a1f11a9a26c89780f5c88c3ddc271669376fff21abcc2d07d52239184b64f8" + }, + { + "path": "0/default/ocdbt.process_0/d/48a739ed7e454a687ff8054229107ff4", + "size_bytes": 8005, + "sha256": "c0b05ca0985b49bb0f697b2aa493d1be94d930abd3cad3296731fe8a2b99e638" + }, + { + "path": "0/default/ocdbt.process_0/d/699c53e18830c1c5729371cfe7992bdf", + "size_bytes": 1696, + "sha256": "ab12000b12e5e22b821bf4c8b3a2aa02d3c81be3769f04b490b74edfcd56ba21" + }, + { + "path": "0/default/ocdbt.process_0/d/9b39f2f8f2d69fd40c8c5fce5b7ef608", + "size_bytes": 237, + "sha256": "90ebd7ae65d5f3140ee1e442b560bc37097fa3c3b7837ee7c7b33de65cdce5e7" + }, + { + "path": "0/default/ocdbt.process_0/d/a721199592d521471bdf14a4837a85a9", + "size_bytes": 13529, + "sha256": "0390c61865b61d220ee41f1df5cd81d9f8292e6438a2b5e4d21a4166c1ee7857" + }, + { + "path": "0/default/ocdbt.process_0/d/f6f509b14bfe5f51f56c22ccb3910b58", + "size_bytes": 386842624, + "sha256": "ef4e67bbd789f279e2b611f79d7b907cd50ae01fdc21d1df5973d88943445580" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 363, + "sha256": "85ce69e56f2d9d14bde689c8b4de70b8a0f381ec19037f9410a7e60cce460e14" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "ed917bbfb993d05c2999d2452a0f027944510a069bcdebf045de9da1b40ad639" + } + ] + }, + { + "id": "ppo100", + "label": "PPO100", + "path": "models/ppo100", + "state_root_sha256": "f840a648a0208577c15247f1119a548580814ce702bf0c1f9bfa41923cf684cb", + "training_step": 44000, + "weights": "ema", + "size_bytes": 580121486, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "c1f7ddf5005fb03561feed2d16be137e542752c773f643d912a244471d5f0071" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/0300e5ac4b49bffe5719d9cfb4762d1a", + "size_bytes": 112892, + "sha256": "7977f2ff2beab7ac7518f82248b59357e186ffc6f850fb6fa42d86cd8d558803" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "20ab214418e3b10cc607ae0e474fc59d0fc55b3c64a9284603eff053ff96b30e" + }, + { + "path": "0/default/ocdbt.process_0/d/3c2bd48694a308b3376ff141afe1a568", + "size_bytes": 428994560, + "sha256": "d9e8752f34ac003dfee90a34c8d6e4430736e00361ed0516c58eeeee5a133505" + }, + { + "path": "0/default/ocdbt.process_0/d/417b2d123a9368637e34a59cad750a06", + "size_bytes": 2060, + "sha256": "fdb904682fae02b2411046985abf38e744b0d1030392584875bb6b44c4043019" + }, + { + "path": "0/default/ocdbt.process_0/d/8dc8e97ab9734180b7d01ab555af1a32", + "size_bytes": 186, + "sha256": "78f1963429d958c839b0072dea545e0c9e8bc3bbf61e0627b723f9f497d79569" + }, + { + "path": "0/default/ocdbt.process_0/d/b50d4837ed106bcce6325d46d7b977ea", + "size_bytes": 1998, + "sha256": "a03930149c9aa0ac2d496dec3781d54833f2ca218d1f8fdea69fade69c170877" + }, + { + "path": "0/default/ocdbt.process_0/d/c08f0db736a2a786d6c2b6a5f8be193d", + "size_bytes": 150392832, + "sha256": "f4cbecf6daf77edfaf04070d82012ae964161961732ccb62b898e89b40814383" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 318, + "sha256": "e2fb22b19dae16a8fa1e57e146c17982d17ffcab4ba6227fdad9eef9a9a0d4f8" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "0cb4ace8d67de9757e41e78b78f6dc604389eaebae255327fe91f36f0cade942" + } + ] + }, + { + "id": "recovery1", + "label": "Recovery1", + "path": "models/recovery1", + "state_root_sha256": "61986cc707f8c05d1d116b57bc5d8d544061e2583f59d084a87764b2231931fb", + "training_step": 37000, + "weights": "ema", + "size_bytes": 580025095, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "6f4011b06330c3e3b3f3a604a302375f2b2cfd89858d595ec4f9b1c21cfe4310" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/d892c60f4654593696558e2196fb3f0f", + "size_bytes": 112929, + "sha256": "79dedd84deeda7d7077e1779f02edcb33a0b7d512257f23ad4f363dc6e0be8cb" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "30f0ee44a34ffb4b44ea90384f9209abaef714e682f94d5e14224a553e664423" + }, + { + "path": "0/default/ocdbt.process_0/d/30c5f0e29e4f4e5045850e4ae8ac8c14", + "size_bytes": 447942656, + "sha256": "014e464927e5fa53f01584e5b7d59c0390f00060e942334aa49c873e4b0bfe88" + }, + { + "path": "0/default/ocdbt.process_0/d/5918aa7c6bcacef5d5a783faa7189dec", + "size_bytes": 562, + "sha256": "64df70ac7fe463a6888793194838db5565073dab4a1f923b4c08b6414f08c9c4" + }, + { + "path": "0/default/ocdbt.process_0/d/7c0764d7c8f8db23a57c6d278af3966f", + "size_bytes": 10940416, + "sha256": "8d3d183c041c89ec2b1098ea2a689ab3ba56547c1c57eba9bcaf11f1d51db15f" + }, + { + "path": "0/default/ocdbt.process_0/d/ae2fbcfef5c278ec27aa5ca43aa085ad", + "size_bytes": 120410112, + "sha256": "266f1e1e48be0aba5fd10fc7496a8dd9f9c101a6d33fa33cf066ea14555b45d8" + }, + { + "path": "0/default/ocdbt.process_0/d/c54c81baeec72d47ab8c23e6dfcd33d6", + "size_bytes": 1460, + "sha256": "c0bcc7f0a6ed14452eb56a59322aecb84cefdc6e46ab4f5f031e4f7cfb921848" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 320, + "sha256": "57d3e5adb986462905ba2c232e166dd60193fe43985f6ca2693015017466828d" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "f77d08b5f5cf3f401f8c9ed8120bb817eee5f1eda45b45869180930e5baacd22" + } + ] + }, + { + "id": "recovery1-transferred-aim500", + "label": "Recovery1 + transferred AIM500", + "path": "models/recovery1-transferred-aim500", + "state_root_sha256": "eab772f2e92d4934f774f6fe36e5fd79532e9c495c88f1bbafb730c998aedb49", + "training_step": 42000, + "weights": "ema", + "size_bytes": 580118655, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "f7ef0b4eed6814e1108310f7e4df01614452e3d220e414fc7e01d7fdb2acf141" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/fc2640fe3dc160b15796bb4e93702707", + "size_bytes": 112894, + "sha256": "50c61d1ec7f080c8f58caf9aea5ea09c286abefca47ac44d49212f753d620a1e" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "2f4edaa35bc63c1ec20c496caa2c4bc4428042cd901363f93b88458a8491ec26" + }, + { + "path": "0/default/ocdbt.process_0/d/262859f5412e92128ec9d77f68d5849e", + "size_bytes": 116887552, + "sha256": "5f79dc226cd0a885d34007b5ae814e43a64f14ccfa688e2b46bb32999c0fff26" + }, + { + "path": "0/default/ocdbt.process_0/d/362947cf928ae683427fe0c00bfe6855", + "size_bytes": 462483456, + "sha256": "6c10e570deef510408653dec66f5575b2f2faafb9e7ff5e63f2351fbf9e31c2f" + }, + { + "path": "0/default/ocdbt.process_0/d/506a3289e41cfc477618955ee791c8b0", + "size_bytes": 186, + "sha256": "78f1963429d958c839b0072dea545e0c9e8bc3bbf61e0627b723f9f497d79569" + }, + { + "path": "0/default/ocdbt.process_0/d/57b7895daa6a39948024d1eadfaa6e87", + "size_bytes": 1998, + "sha256": "a03930149c9aa0ac2d496dec3781d54833f2ca218d1f8fdea69fade69c170877" + }, + { + "path": "0/default/ocdbt.process_0/d/9a0dde52525302e7d178310d631e3cc7", + "size_bytes": 13516, + "sha256": "f128e89301603e200b7a6c3a00c64e3874dfe98b258f0552b98b36aed82bd9bc" + }, + { + "path": "0/default/ocdbt.process_0/d/b48ef8617c1cec65c1ef28d3a634acd3", + "size_bytes": 2060, + "sha256": "fdb904682fae02b2411046985abf38e744b0d1030392584875bb6b44c4043019" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 353, + "sha256": "4a55b7aa0ff3d4f23a4f4be3a740b4fb13fdc2703d3c545752c5ce9fd68abdb1" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "46b8102cfc074c4bcae2c2252ddd13b86375aab1072610946633f8a86011542b" + } + ] + }, + { + "id": "recovery100", + "label": "Recovery100", + "path": "models/recovery100", + "state_root_sha256": "932ab1e9013d756c3a8c80c1c7c3d9c6df2f6bc46426733805fabff9d1839460", + "training_step": 40000, + "weights": "ema", + "size_bytes": 580051815, + "files": [ + { + "path": "0/_CHECKPOINT_METADATA", + "size_bytes": 275, + "sha256": "736cd91bf2216f2f434d5b7dd46ebdd41814a23d03d4a159efdb2962c274839f" + }, + { + "path": "0/default/_METADATA", + "size_bytes": 234269, + "sha256": "640a3dae19ca05967a57f5a33f80a21045d05995b55daaaef4097997aa5e5789" + }, + { + "path": "0/default/_sharding", + "size_bytes": 166020, + "sha256": "beffe5eb35c116e9c2275818a8278d2bf1e7311acceb3b27977d605e07548c6d" + }, + { + "path": "0/default/array_metadatas/process_0", + "size_bytes": 85571, + "sha256": "110bcbcd547cb533982a4792e35531d5f7ad8197edd9ac709fe55060170ba07b" + }, + { + "path": "0/default/d/83878ea0eec4a9c18daeb44b066196fb", + "size_bytes": 112810, + "sha256": "cb3f3ae6412a95e2c6ca942b722c09a20a7f89c31f816dd0abca5234e87092ca" + }, + { + "path": "0/default/manifest.ocdbt", + "size_bytes": 120, + "sha256": "7fb842e3e6798d4984c642836aaa019bc649368eb6e9d222474f44055abdfaff" + }, + { + "path": "0/default/ocdbt.process_0/d/4d16686f37d7842d645df7bc0fbc9e78", + "size_bytes": 199, + "sha256": "77e6f800feb9f959ba9c8c693b4f6f183a9c090a693f4359f6381f405287634f" + }, + { + "path": "0/default/ocdbt.process_0/d/4f2e215ac613b852f54f8856247bbadc", + "size_bytes": 2096, + "sha256": "f27f5213b11e8447c6a86382042f9c96bdc59f3a907f6e066dd7f525e9deab74" + }, + { + "path": "0/default/ocdbt.process_0/d/77cebf3655ec723496319cdb1c72edbc", + "size_bytes": 402022400, + "sha256": "debe6d334fbfd88b24c3f1bf0eb8445c2d5ffdd63cc8f7f72a3a8536441189d4" + }, + { + "path": "0/default/ocdbt.process_0/d/989b03147ff1281329a5cc0adfa3cea3", + "size_bytes": 1998, + "sha256": "a03930149c9aa0ac2d496dec3781d54833f2ca218d1f8fdea69fade69c170877" + }, + { + "path": "0/default/ocdbt.process_0/d/d1816013dc8ccef6782787251339f26d", + "size_bytes": 177295360, + "sha256": "7c6b80c42ae19b2a8a4fff6c8e85fca1200ba472d247c252316988106af97558" + }, + { + "path": "0/default/ocdbt.process_0/manifest.ocdbt", + "size_bytes": 312, + "sha256": "88f561a6c5808e5ff9e69114d7cb60ca3fafa760668b37e6da9df03ceebc315d" + }, + { + "path": "checkpoint_manifest.json", + "size_bytes": 130385, + "sha256": "c39dd879a1f35b419278ad766dcb89ffe7a707fc45691dc46fb308b65700d222" + } + ] } ] } diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/_CHECKPOINT_METADATA b/models/ppo1-aim500-fixed-reward-estimate/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..f41fd31b9e58f39b679674d5755c49e7ed3c5bd5 --- /dev/null +++ b/models/ppo1-aim500-fixed-reward-estimate/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786293418670830991, "commit_timestamp_nsecs": 1786293419315247467, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/_METADATA b/models/ppo1-aim500-fixed-reward-estimate/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-aim500-fixed-reward-estimate/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/_sharding b/models/ppo1-aim500-fixed-reward-estimate/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-aim500-fixed-reward-estimate/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/array_metadatas/process_0 b/models/ppo1-aim500-fixed-reward-estimate/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-aim500-fixed-reward-estimate/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/d/15a85f646bafb55f51a18966a8a4eba3 b/models/ppo1-aim500-fixed-reward-estimate/0/default/d/15a85f646bafb55f51a18966a8a4eba3 new file mode 100644 index 0000000000000000000000000000000000000000..5ee5a52cf0427c40c8ab3df96bdc4c218ff7058e --- /dev/null +++ b/models/ppo1-aim500-fixed-reward-estimate/0/default/d/15a85f646bafb55f51a18966a8a4eba3 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6b16cd6917850e5d2fb161c7b94ea3fdd4b542bd928f05e96b1884a575c04a4c +size 112946 diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/manifest.ocdbt b/models/ppo1-aim500-fixed-reward-estimate/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..256932a4aff6eec711e837ca639cbda44c76e7b7 Binary files /dev/null and b/models/ppo1-aim500-fixed-reward-estimate/0/default/manifest.ocdbt differ diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/1691c69e471866bf5d174ba00029a7d2 b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/1691c69e471866bf5d174ba00029a7d2 new file mode 100644 index 0000000000000000000000000000000000000000..a29d34f72700ac987b417e360f4ab797e7ba77c1 --- /dev/null +++ b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/1691c69e471866bf5d174ba00029a7d2 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5e9efd37754a3b542b24415f98c2355e3cada3c6da515eaffb3e31514be9ea90 +size 541900800 diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/6cdfe6d8112774ebea68968a8b8b2450 b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/6cdfe6d8112774ebea68968a8b8b2450 new file mode 100644 index 0000000000000000000000000000000000000000..24bd3a33e71001803fceab3b4d2038a2d8fcbc17 Binary files /dev/null and b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/6cdfe6d8112774ebea68968a8b8b2450 differ diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/7931d26a51fcf24844ca857b5b1775cb b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/7931d26a51fcf24844ca857b5b1775cb new file mode 100644 index 0000000000000000000000000000000000000000..d41054bd844c5633a6580b9ade5a3aa36a7f1e1a Binary files /dev/null and b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/7931d26a51fcf24844ca857b5b1775cb differ diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/a5c70bae84187efd011e1c2d1ee4be1e b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/a5c70bae84187efd011e1c2d1ee4be1e new file mode 100644 index 0000000000000000000000000000000000000000..eb5738dadfa525c21d5b96d440926eb9bfd256c5 Binary files /dev/null and b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/a5c70bae84187efd011e1c2d1ee4be1e differ diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/f2e7e5c758eb4773c9cee341213d7c29 b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/f2e7e5c758eb4773c9cee341213d7c29 new file mode 100644 index 0000000000000000000000000000000000000000..1f0bd4c2f68109b32e1db25449aca60819766996 --- /dev/null +++ b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/d/f2e7e5c758eb4773c9cee341213d7c29 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f3c66b0ca613fd2df0d7838e15033b8f2bf4e0e2dc05143d467b730c202610ff +size 37326848 diff --git a/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..f5a14e5e4dcd1031849fd4a40078fc578bd6f127 Binary files /dev/null and b/models/ppo1-aim500-fixed-reward-estimate/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-aim500-fixed-reward-estimate/checkpoint_manifest.json b/models/ppo1-aim500-fixed-reward-estimate/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..be54c2faed9dfa90c2cac8aa45b0e95088511d12 --- /dev/null +++ b/models/ppo1-aim500-fixed-reward-estimate/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 42000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "95e20766bdc2068c26c115ee32cb8e8a8a4e6bc00b498f083ceba89ca8be7dca", + "runtime_config_sha256": "95e20766bdc2068c26c115ee32cb8e8a8a4e6bc00b498f083ceba89ca8be7dca", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "95e20766bdc2068c26c115ee32cb8e8a8a4e6bc00b498f083ceba89ca8be7dca", + "schema_version": 1, + "training_contract_sha256": "bfd84af91ef379de7b3ee6f992a5fad5501dfca61b30b3648de5f800d32450e3" + }, + "sha256": "0b4d2a1e21f203b20b4313f25c9a2bd1840f589365a09ad0f9d2126c2aadaff1" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "0e3eb5dca51f099976ae56d5226ad35c8b78073638458158808f8446704732f1", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "6fbaedbf7fa89bc5c5c00732eaa64a78158cbeb5573f5434fbb9700d2f6a8fc5", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "d6afe68a6cd0fc81cfd9c1cd18e171866ed070bc6d4730d2d098d25f4bbc6e41", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "9ffcde10d70deeebef597dd2b4fcfa63af583f1b9d11a5d5b3127fa796d1d3b0", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "acf5d2b207d34ba5cbb6102f1d73aa3014110308c20f202234d5be4824dc9e3b", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "32741774d268c8c646d9941fc5823397cf95058131e9fb19c6fad6da849cf8c5", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "d462721905f7d0ac17b04e1ae93f7bf6efe34489b1c85d771ed2f6e836876621", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "05383f6a2b54c0c6215043f6accfc972154bd079caa6370ba60c7a7ebc3c9de6", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "ecdf75a92b08a35d8a40eff55c51fa6a98a79761a93a18de8bbe4d80f72affdc", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "d2141b89dd7f674d3cc359fc006d74392d1deba43ccb4555fae6234eb3b936d8", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "6c9211d0ce8bd4cd41bad3bb550bf740365f8fdf54043bfaa172cf70b3bf726b", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "4c7d8cd4366bba7b97a7f1938c0ceeb50bdbe32a6cb8c72b473370ac24d020ac", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "f15f51b21d4581e97021a11a5119bed9a274f56e8ce8bfb84aab4158c236da99", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "527e9be1fc99a3f21e2cf8b88bb8dff60491d96a20f0009f3db0c4cf2dbbf4da", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "f6d10f765c1cbe6b5030152e815b25803d7a3d661ea646615cfc10e169b47c34", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "8d2c83bc44c61df528c464766633b83fef94aa50839ec356eb326732d29282f8", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "bf3d9d0a04d0c3357b28a18bb4db01d128d09ef65ce7866462e5d3aba5ef3328", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "3b3cfc410b279c03e84d8a6bdda9227d9659c1f54874f05362f20a3234c992d2", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "747c9c83e81ff448ee2de1628cbc8bb31f77548cc77d476727601d9ae17d7624", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "052ac175e767ab36acbe130a136e1511c85cd9d157ed9186933c21fea7e4bb56", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "5577eb5c4a65efe4129c23b03cf8f3e6dcc0dfe82883b7fb133c376d9ce05efb", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "ddbd8c6cbae5cdf2130e92698d6327bcd27ebf7e26e5f00e9b2ad00df3136a37", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "0f83a19c65a4a9e9aadf0fe61f77568bd093fbc35c6a259b1ba028c1d625fe13", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "ee34c9582aca3f537f51166a96edfa2028e00b1dbb0f2bc85b209e970b7839a9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "92aff82fad4eaafd4131c028e3992c98c4f25952d28b75075584a5e3de69605b", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "0aaf5e52ae2aebcef158ae4858203b92582151e8920cf3329b276b6d72f3ff31", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "273fe833a81c40743da7759cc1165586b2162a60a7bee488413a1f5efbecbb7a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "314ff313c86daa29ec09071a20487f365160042d673da540a964577b485d1be4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "67ceb450cc9c2a7f92eaf4589c31ae54f16eeb7575bdf302fc1d4bfe253d2748", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b23b8e691ea8d65c147c0effbe5275f65b3fbd9766e5ee45bce4f36e31d05f78", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "c8539d359755a932852a0e976640b15cf891fea7374e97b8edefbdc2039ec46d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "65b699805d64bbb5deaab5b635e90bac11e1ea2557a61ad9c6bcf4ee483a45e1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e2f463ba1ef7edc4b8b71fee1de774d969a9d9c7cdec16ed6c48c92ea183eb2f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "6aabc4ce1c53a5bc0943d2f5ed3f68a50cd5fba20e31e16ec20b454b18742f4a", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d90ab368dd5df64a03427b4a070d8abfe0d1b84c92cc9599df268270cd67d6b0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "45ab9d030cf2a5d1df23fe6b45ae4b77e90f76a9e90360a6fb957b3ed534c249", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "ba0c4cae560c45e9be1ff83a503602d49cb11de6b4a76786ff8c32fa9fc7d973", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "647fd50be1cc84b1acb8df532ad43b5b7068a90606bba97e3c0a170b5622bdab", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "18e6681837e98c465b820ef3173772f58415eaa8db70ede8a3736337120c419c", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "c032abfd3d906c7c58c007417776b3ba6c7768c20bdb50d793a15a7ceeefc831", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "383e0906b90c39a679ffdbd7316963b8c84597ce64d76f7c42fc429e6fd60f39", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "3a65fd219d221dc1fae5c7ef0ccdf761ea6b2ad083cfebc86ed83f10da62427a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "8c461e09212055fa4c0277b8ae1e3c984b87f3dcefd6d7cb882e08dc825e1860", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "3be655f4ff19371befd21d6f753ec6e478417898f4d446cf7ee404ea54f9aa19", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "06e32da2524804a1a55d59045ef32cea1689dd39e0d7d89e355c33dd54b659a8", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "32a0b57d76a910363a20dd5a1d241910060bf20169a9fe8db956d7b698797060", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "bcc2abeeb509c4aff0715386a5297610733bd6b5d5d06cf1b8d9791338b5226c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "0622cb216738bcc4f05fe1ddaf7ca6caf418d8d87983d6c7fc501cd9e683f2e5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "e1b1d3930912f9eadd9a59b59503778ccf21934193f7f1c3d1dd4e9bf8e8e4bd", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ef46249b249d84582e992ef06729a38898400b69d742a932cb0a996006ce69f0", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "82cf14cdb988b7bc882667b1fe81785cf8d0317636a29dfa5fca91a54ee6c8a9", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "72fcd1886134c9d213e2e52591ce0462180dd1937d80e4b2a57c6e6f7fd2cb55", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "30f1e06ee1d5009fbe0dbb46e09ef8b19fab58c0c8004e0380c64e395f09e77f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "88cb62e653769163e50f3604c1773beb3fc6551b5790ed734479b9025d50fc7c", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c3b6de4deb4a76e6853575994233768f9cd481df042764d9d15e06a22a4c5854", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "4931324fc099f092d4d959785bba41c94a3df39c32b0a3ff21bd8742144c43ea", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "a4f42a192aa9f507d7e6a5554e63bf217f4e3d56f59242607876073249bce475", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "29d4668d12b58428440f36f4dd7ea926733e19397de01df601ae5b301a452fd4", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "34edfc9a49d785e2a9b6a48ff047f0b72ed9a6cff055f1195228fa8a20f9015b", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "6ebb10df1450f7a679ea9c48935d0cc4f9fb19238ad60310f5816aeb601513c0", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "9837e5a7cd50c9e5f54c3bbd79dcc0de5e94dba34bad2da28a55a5628357ac2e", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "47788b3ac733ee87dc9f2be4af9cefc7d96c8665498a2d1eb21ae33c5c7319b2", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "9a914cf417b2adca37260b135e08d91ca70c54e99d8d28eb93daabfb28f1831d", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "b33afa9c3598ee8c4a8318c1cf282bd25b403f301ff38a0b5dbb337f505e5aca", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5c94c5e25f671fed11ada94698500e09a49250fa27973d7c2ee434cb22a68122", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "5c91a09f37094b4193814c6022f98366bd3c6931ab1d325b45f025d417d9d1c8", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "81334e289684938ad799276dd41c07ffafc7842af2352563ff19e2087cf6c0ae", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "10d45e8ae4569cc6af57766be8966678f324e28c3058c2db2be5bec9e420b293", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "58ef36fc961201ea8ca10568208cab9fbae84f0919526a45a4b989f3a7bcf16d", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "14bc6d46e0e8f8d6daa0348efad5133b821faafcf2fdbb9bf38dbf09df82c2fe", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "ba50d849c4705f732b2cd90611a1eebb743687e1f52269c50d08fbaaa76dc189", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "c4655705e472a8c94eaddd7bf89d1dcb77fcd00a0b6134ed1081b3297bd57e62", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "7e2e46ebf93618639d4657b063437ce57adb107e4dbde9f88529db0566585068", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "442142b8900cc9be02a9515635598b0954415b51e3d4fbf125b6aa1123807885", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "1c75ea1d506f7b5ef9fdf543e1021417edcf1191cdd3637d2169717ba48f1184", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "361929cab4427a368b6ec871f18edb4ce567618a7e88470686014a37b4d284e5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "0854422c6ab6059e6e4de05f72aa800562e5a6c0b6a5359bf430aa9c49bf3def", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "deebb815232e377df473ce56828a36f85cb7e0558740d78143655a60a6a39e89", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "634a6266665dc1546023cea60c7b4d005fc2220b10aa0025cf3cc0b81ae21b9e", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "ed3ccb0484fcde7f429262fbcccfe89eefc512ba14d81b62bccae2771fa084cf", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "573e487aff7efbe7d525f9cd0dfd4ff5086dd5a38506d842b93756171cf06666", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "3010b3f6e9468e0ca25d771becd7c79921cdec417f05e719dea6882f6fdb00a6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "18b81cdcdb4c8d7222ebe0bd35832a3a3dc99a1495ca06e9407fc3892a8a264c", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "80332abf1bc4e59e0e5eb2738bc324b53acfc4f72fd78e9fd3151653b4db1c3a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "1ce0f593389a425ceb9ab60d31ee96efb51c5355d8ded2f279e74470260b70d2", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "7da76670fe13143b5013accbcb1a4d8d634b4962eddfe9e5027b8c5903b651fb", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "4d8cf306c1b81c0d68246ff9ece6c88a496780c02a33fad502cf9d5bfc006397", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "4b978787ea47e6b532ecc62f3e3106e896cf02821a0879343787137a8236cadd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "71fbd64efff2c2f7aca51454f89ef25a8dee15f1ec101dd6da1bc5322dd4af6b", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "4b113397d8f7f5e7d3d479bf1ce51fcd143bfd660f567b70e9e142d75ba15849", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "b53b07a019ba1959b3d06edfb343387b569a9af28e462e57d0f5f04ebac4e5e5", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "43986a38d7a8714138956cd54c0ba9c12cb432cdc852bdd1c0d11f26a2be3008", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "c31b83604bdc4f394791457879e00098dea8b18a7fbb9b31f5b017c1ee9cb0c4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e15004c03170b7e04fb4a291446baf5549478c5c7e23c8dc26cb1992dc88c066", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "1790b69f30fa188c91dfd3cf6439222b2acebe0e16f72d1c1b0930c1343cd8a3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "ec3952a1d3f1de19b2f095dcb78a58c8cfc6317b0151abda7df57c40f56ef7fb", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "469d68096b3908433bd4e7af8b23d10432be67ac83eeedc158c183abfd638d30", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "2cd2cfb31201468f02a75a13fe62ebb77c28a62a62a0d39b83cf839fed6e5b4f", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "f0231de6fa3a40d437c331e546931cbcfb83953fbc010be7d18d61a1fa7da8ea", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "144ae6613698ed2c1b7e68c4ffd0606d6eb413efefc996d3eae5fc11852d66b5", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d96da959b8d60093c8f98d2df97bd5b47086ddf807cfc5797d6b63074554abe6", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "3c641801b7e8ac85a14e6d12bf61002df6361711a91a7fd3f6816f3176bc66a6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "dc170d7df6b21a8a648fb16ce66dfc09e0bf0a58fc8bf9a9352aa0bfd1d36cd1", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "21e63485e448ce3d9a8664346ac2b2a9ac2e21dc78eb22737e47f0addc391b83", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "1d2764692ecdd53cc6f3932eca97da92fb5cec452a2d64ee65c2fac37771b406", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "cba62be3f00aa3d899452ef8c25a6acbce9349034de08a4e15b1740d046fd16b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "8896dbdbc8360943a2bbb9ef040e5d5ffcd0fbb52754a999518fadb562a97dcf", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e50d7dd4388e244c3a0dec2fb18ffce07ee587f50769fc7f99125aeae7c2b47c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "2348b2fb4047fd1ee05a3b8b6960ae881fd6db536cc95a8262eae8780871ff02", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "10b25d9494ba6dd423fe6eab196548c109bcdcf3463aad47a12e65da6ac41c97", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "518f734786e89bc304aae46f7b477761701dacd3178c0b43a99ab9b34c079000", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "70707955f237d5e37e5a31d9d04c9f0cd8afa3d5d8d1f618a7df6ec092253568", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "df241362e4e4ff6b1ddb5bbaf141c8700445793d15ee14465cf6318c217de203", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "1a879064c80ba7b6af90f7ded2885bb6b22636d9d45e7a4f419655fee0fa841c", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "395b34d48743ba70f2ce510c1240d5e213bc3f7137dfb0587fe3c485c59cb893", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "788d50e17f0e14cfb790cd617ab9571ae7eb4755163474df5b1a1f9d068dadf8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "754283885f90b736003ec2d680a12bb9c5aa453a601dc561fc948c19eab44d2a", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "91ecabd14393d1ab0fa5277700930513988f636b0ad48e9a8b74618bc4e11b5a", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "7d8618a51a5fc61c963befb1005af999633a0e2ff507fc1c769ec3f46c29e645", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "2bf4fb4fbf8ba9d6cecc1cc8ad9c207ff38afbb5374aef58d7a97701b0b28ca6", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "de43ad92bdab6bddc287c7ff4167257740d81c8760fa5a347924c257e513fc56", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "c6e485000aa0e4865026f97e9eaae14f6e7b37e51dcfef88fad1f068454e0206", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "ce4fd5de06eae4beb9ac05aa75e236530580fbc42da2ef8374cbddfb1ebac09c", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "5adf68e698fcbf32ad90fecb575a54d326328ddcad77546792e7e9f8f6ac6f98", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "574f0d7338e45e12e521e31ebda640621f8ad5b6fc61f02579fbf73bd2fa6e2a", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "bd8cd66e9397a1ebb1df66c6cd50f746500a3a628676875cf9a5d1f680d7db11", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "e30dbbac586776d424f47791ec97ddab5d38b2ae73d9508c3448f966fa3fa6e9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "a68e0ff80647c616fd00acc691ef4556601fb132c80c8ed1172f1c234609e0cc", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "9988ec142f34703acf375bb97e7184da0952d5867240fb61692e725e53c4b138", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "95e261f27322118f4281c8373069e5825a0f7d44e685de24ce774d1f2c355b20", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "730ac3464b6a7e31408d422ffe82939336e7375f4c16a4a8e45a19d4fabda78b", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "94ab8ea4c4e751a8ffc6617ee7cf0c5633cf124e6cd83f27ec0cbfd9d22ab742", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "3ccaf630b973b4dda8d81f59270eb79ddf261e755b7c854b4ec1b451ec075924", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "993ece8287154b797554ed49c3870d533f71b5743335b79d1c953a963bb691f3", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "f906f92ecfc4f75e4ebc3ecfb3fe758e54c762f692ee0a35d7342a6fb5a70729", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "ca3f7642895e17e5f5dcd7480c298f62411e4ac9818e124ccd9543afa675306e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "377220cb77ebd763d2f54a4f7ad035f471e80abb64f7e0263cb0aab07ff1bd9f", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "d1d0db8471493a1604954dc2e0023b02d2108ffe4b00e4f2ff3622d31bfc9e59", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "dd808fc724d151769eb4d4f6a7325df12fd62cc8b47a60e2361bf4c1a5047186", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "f9422f1614bff6a2e8687b08a3ef827d1167877881d8a536b592a85c49e7c53a", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d00843ffa9d726dc3126cf056711b078b37fc582dd60731e4c6cdc5b0148079f", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6995b17f9e224ea60338b426c01e7dbe9f4b82119a3ba613bc95bd481fc06e48", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "1cc6b276ca8eed8f74b1594a427c9e4c5656fe632976c966f8fbdc501acc67cb", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "9b1b981610dd0ae4bbd57a256fa157613cec333836580c14a326a09562f1ae5c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "59226f886c04fab653d499e78c53448199297225071d8ff189a66f500e8bf356", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "ae0760341e4c31d6f7f1150d205c61344c8b70de0f6b9aa608ba5f9b8e06285c", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "9b7b7d2a7684867d35dfce120f80015223558e9084e6559b7e97260a00e9640a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8ca599a862220cc86834931e346e36f97a5f624ed4b3515f3ecde763ec03d82a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "27413d268c78c12412750dfcb795eda2529e9240ed7c7e847a55fba100a652a2", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "fb7ccd23feb1c7520e8255c11a52c1b66d6883c4f2e3af1c8a38578b52648451", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "96fc25b3827a031a38ddd30e07c588423281f457c4e0a1d61d62f76efe01991f", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "80305dc887e284f0887988dcb23701a063056b281088f2abb082998e7a8baf4b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "64a1dd720ff4c741761188d21dfdd859bfdd3796487c0eeb877d9cac4c676737", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "5fcac4f8857b3d48a8cee8481c4d6c6a5842b74cebf482f5cfb5d9e9558676c4", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "55af811186b4270458a02208bb4c868b2e3d8533a688c55d996a9e53a43f916b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "cb991eba6a1c1abb7548dd48799e3668283f54272fd6414fa3b48618d9c657bf", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "13f38373b7886e5cc4b7cacbf85728fa15335f7c6cec410c27d32b2f930328c6", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "5af7339d671eaf8a0a3de145a38bcd24e4ad4fb03934010fd04b62164f9ab973", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "2e49d3af9155f9f1bfead372b17876bf6826f36420831aaa5db71ef41adc1fa2", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "009ff7f3d1a8a781cea871e252f0125ff28222e4de11d837473f29d0eb030f04", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "3087cfb5107fbac119f0429fb4619500ed2033c20988c4c19af70169abcb7b7f", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "0033c1fb25a384cfe806d7443ce669356463196e9c04158912b850b06d405680", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "532c9fae7f3a1e4b4afbb7109d2f86e57785c201cac5b6ba65c4414d51384eae", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "3c771d28b8b187f90e50367363ac5c49336dc8bb6f7bf8f87e8e7ebb09ab8011", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "9569b82707e93ae71f7006d58d6c65662c0f4669aadcf9aa9a71dead9cfa847e", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "27d799bf0f300fc2925b84e056784bebaf7e5c5612ec760a45e6dd053aaa3920", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "5fcffa154fbbb2d93f19e002304e45390e4909ad8252c1142941ee02e8e0611d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "4721dc0ee6d6ac09ab98269f118a3bb86ed14858f71769088cf1ccefb0465cab", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "5815a9f5edabeeb238bff385937b4f6efe0709f4ab1c15d9510e23c186476f69", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "8c0326b6ee5c6d7552a544a31c9f02a5e7741f23e4af7bc2a14a90806cbb44c1", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "680097c377861274cc74da76e561924a49d5daf0b6ec535fa67b78dce9e1d8f7", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "072c1e362197c91d707facef9982d6b72b38f28e9eb5bd52b1d084f4ccc812f0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "98d68cc25071d56cc9fa4efeaa640de804bc2d8d31fa07a5cb4c89f61e1a7c63", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "a035437d9f9ca4baec73935677cd78980d1b00d8835a1e8ebd76ff4d13a38a4e", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "54a3544c4a4d20ad44bdfa5099426ac66d5f3c9e2c6c624e74451f3f64835564", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "5367af128a93b0d1802af15f89573dd97b4b6bcaeff6aef557b72235d96bc6a9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "7c703f6e3e7225d07e19ec188705fce3016d2b5809cd003cefc384352b83fa1e", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "9cadad2eb03f5de4c503222882607e3af7152a176e86a34504c78d65cd39b4f5", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "2f74c66469765358569d7e9e0527d9e5dea595382e56c4832689bb1b6fea4959", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "0ab7f5aeb5fb0b398598feca0d7de115607ce274cf7f93a945ddd7893b3a1b04", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f96b5a15cbda0c14218f9d0275604d0ca941d73997383a431d546f963c29c007", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "04c5f4ab98ad025f6018b7854e77b65536e55be5c2f7ae165da2a7af5986941a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "5503a50a360535ec44bc2735d9eee855408f08321afb63591cef6f5b040c28e9", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "86dc275da2b3773ddaf1f2a32da5d4ca7fc39916ea2abfa3a95d260095d39329", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "81d0eb66e52b949e0f29a24b2dbc0134043fa13c64160de8a10b09cceb842e8b", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "bf1ccec051a145d9625db4b8d461b46ac1360bc4953803d8fcdb3665f4f5f8ca", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6a34d0179345bc489eddab0bfed85dd561ec7678d9be5ce821bd159fa1d19630", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "2896befd87c3222df831dfb85b8bc07b44658cda9ed819257b1503732e42f8c9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "a95bdb024065908478ccfa29f5c80a03086f313a4a69f9b64865b5be5d987e8d", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e8b6f05b9ec04c5725f93a0f98d70a17dcb678cd3c07a521914b25b1ff2c8c63", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "2a23d8b1a7669414f869acf9785175f0c94980d33ceb151ad557b7a5144b1f45", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "3891880f3a81d936fdc7184e30952a3c5e251375c8fd80fc3d667e7e7a7e2751", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "16073af81a5388dbe604c4e37b762b40e2132d00126d8067a9b9a95e9623146c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "86b208df006f851c912f1b88df1228e41381708dd158e49dfdbf3c950aa07d48", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d881c6dab87f43661a83b7bb0c618b117bad4ba8804a886ec96e1f7cb37ce6d0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "d70817a7bde2229bb1903fb54f2fbcd61a52426f63130674082d6fd9567277f8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "a5eabf0ec0d65f081063def9c5cb235e4f12dbcc95aaaa9d7d10080b37459947", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "dd65466843ce49658607c28abf7ccbb213df31093fd43daca876e3c204ea148c", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "c1f048dfead32cf397278a7e9be189f992981765d56eb9fa29fc7bb71e15fab0", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "4173675b8f7f7700eabd837c9029a2ac1ac9f9ce4b9d84ab20cc11931e0cd4fb", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "8a486ec31c0a424db0641807c024351d18d6400e43e64ed975903fe796aeea24", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "bfeabd18c392c3672f58223ae2e58d6b0d15257e330be2469ad0704a0c627bc9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "68cc4ef1dcb2a8d685f32c2f032b4f0de87b5db5434a440df0bd0b74f9faca87", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c30c8fcf4f8b047f28384e28c61606e4780a732dcd485d869178e6dcf3f3c19d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "f868271d5b95da12c1cc0ff7781c550b1d0c77fd7ec1e943c328157e1cb0e296", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "85ec193fabcc7392b9c932a9d4262db873d2cf100869aac7d95587a68cedc48a", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6e90ba449aa030c7c4fb21a37bd4411ea04fc2e45ab810c7373e2afaabdec1f4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "bbf7868f23d6d5174011dee89a1c16ccdc2876aef8bb207ab42b9a933add1633", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "bb8e2fd56e507f4938a89e22924be93d88db792233dbe4a2ba0be91e36934630", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "8415ada5a5126961e73b0d885eae8f72fcd8cfeda8ad17bbb71fd6aee76c27a1", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "ae6cf9378ce1a52753f8651844eeafd27247fc32dfebf63926780272b74b5751", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "d0d2214c0d2e1a224d7bcf99047683a38545cec108eabdb9cca177175c269177", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "8a250a9db6cc078615edb6303bb2e864e0fa76e6b1997d963a4a908ae7037119", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "72766b9bfa0786235cf34a271aa0794f58447c1ae7ede5c09992f2a4058f4130", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c046a37d7332fa78550d631231815704ab90db24ab808e878edb9380f37f324d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "1ddacfe25a124a8e537998298380caa50ebf3f16d2575219f2fb5d13b8859466", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "38a74dd55bac5a5fd8b629a81a0bd29241cbf7e94ad892af77a6043702b3450d", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "4af72108d506e7ec52589614a1f343fc3d0a5b08e9516441cf8e4f846e27a7fa", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "d1df2bbb191f3ec994a356be00a4ba81591cc1fbf730d7d30d669ef7a7bf7706", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "f8c79c333ad1aa8f0a7dcc76ae21b836328ee7f6854e5997211e8a5c447950f8", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "dd54575925da29f9ade3d30987d20a4d565b49e2e9c4c443abee0439d0c8a9bf", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "c25fcbdbbf8ac4ff9b00c307bad9229cd395eb8e7eef3a9f26aee9811e161634", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "cac18b11aa2f83c6f43b17b341fe94dddadb541298f826733c884da00056b5ad", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "e5c79eab602011657e29de747ea477c2870595eb94f8b7934017547e5f68d32b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "37dbd260802245fe67cf022a6b0db28a5353723bd6519b51c3c9ff03a0865494", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "288326730e46d31ee63c643c94da738343b7b70e8db094ff87771ca7c5107e2d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a1e8fef0620c597d90215204d61cb826b48c4b74e504a08d9d83c1b0cbd2fa14", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d487ed161017ca1d37da4458283b128206fcc1d40d80097319d50f9c778b3d84", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "dcec260cbf79c9250d4638df8b4f3da226829ef6e6ac71a2373945b18e5324d1", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "9df76caf99c1244cd76cd089a8cfe0c93e000833362307a60b6a0a220a69b70a", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "24df811dac88aa94e268fb325b6d134645031bdbba97b04fe75d991df8d7d2c8", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "aeb53b2f5b2b35e0f601df2c8f1543228784bfbfd9fa6fb27d62bf9cf2ad8631", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d80aea1aa20df68f30fab13bd24f395dc147b59ec528d7326907ab032eff0d00", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e696333367f7939cc101f5f36aa1f5ce73b7f723dd5a47b4a4293f9eb5e0e091", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "698ea9a18f6fa7cff00fe0d83f12923930773253152e5b69b99afb3b0dbf750c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "a9b4b2ca696253b48b13fd60302007aa5f3ebb25eea436a889a95e84497ab968", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "6e6f3c168ded259a6583f463777527bec20da67265eccacadad4746251ac93eb", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "6e6f595d47a3448cbd30c8eab942766a800c5af4e530c0b6a048a92704b0e2b8", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "e3e75a6050c229e952956e2c11a50a5706753ec71d7e8d3a28d89d62ad032c13", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "ed82399979e3d6ebb2bc6965cd7d01c5cc0478ff961c25aaeeb9840a7f5ff56a", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "970617646f60cadb6d5c6efc8b3be6ebbc41c1aa2c2d0c0b5d4e37a9f5edf6c5", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "5762ff4c704f57d73aa2b78e8664f751279414985dcfeafb1d3c2f5c79a960d1", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "d841b0b9247d52d5a2cde02000855fea7a3ce153f42b2d60583bf1b0a1614b56", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "72464211acfab520c145876e5bd1e165d449f305a424ac0ebbf0ac8ce5af3389", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "6e4df0f58193556eaf22ec34bae9314fdb9e2c84e607222d2f66c5d85e954f7e", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f204b39c2b76e7ab7047b6bc61f01a5d6e1e78e2412a7a26077c1eda2f2f6b9f", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "c9b079525927b89eb2f02aa49ba1f21b25716a7586f5df90494e43c8b36c3509", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "29213c5cfa9291f63e0af664909c24ef9e31ee669e58f2b192d8a2db535f9bd5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c5b21cbeebb412bb16f7cd62291499788eb30855a67b41b2e0e52d234095ca58", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "addcbf39835f6e1cb97ad8ccf0a2b80bd7dfa8cc6f52cdc296b40b726da59ede", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "fc1080168ace468bbfa409e73d46694323c11febb57c6384704a1d3e5107da03", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b0c7540553e99b9f9f670bf043d925e27a5518a550732cdc929a6568c9557e33", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "15eb1b9dbaf11bf4018cf891d8be9cd62e1393830ac5426d0491f9b29e672334", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "ad0eeae0a5a00280d54822b1bef33ad33a35ce3aad6a82925613c494ac97ddde", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "9a7293ded2a4fb278e9a61d7bf45e19d59a0dfb8cce5aa01bd064d51028498aa", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "168b2a87cfda853c5cbe8be22fe02dce73ca2399552dedaecbdc1e66384f14eb", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "56150f1455f7aec4082ec7938f2f8161a3aa07320b1285b273113871face7baf", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "24780d3293d74457c42332d4a1f63b2fb3cb18295f6e7796af644826e3d27cbb" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "5975e49b39122fec9983acb172dda693a902ddf574750920b74e56ef8e67f4dc", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 42000, + "training_contract_sha256": "bfd84af91ef379de7b3ee6f992a5fad5501dfca61b30b3648de5f800d32450e3" + } +} diff --git a/models/ppo1-aim500-mi-only/0/_CHECKPOINT_METADATA b/models/ppo1-aim500-mi-only/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..2aa911ff3a14c58686cf952bd83d78841056fddd --- /dev/null +++ b/models/ppo1-aim500-mi-only/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786375424897331670, "commit_timestamp_nsecs": 1786375425581838503, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-aim500-mi-only/0/default/_METADATA b/models/ppo1-aim500-mi-only/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-aim500-mi-only/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-aim500-mi-only/0/default/_sharding b/models/ppo1-aim500-mi-only/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-aim500-mi-only/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-aim500-mi-only/0/default/array_metadatas/process_0 b/models/ppo1-aim500-mi-only/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-aim500-mi-only/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-aim500-mi-only/0/default/d/317d4d2f6c3a4abcade3f68dc576f6a9 b/models/ppo1-aim500-mi-only/0/default/d/317d4d2f6c3a4abcade3f68dc576f6a9 new file mode 100644 index 0000000000000000000000000000000000000000..1abbb4c7b331391a983b43b3923260da7f4b274b --- /dev/null +++ b/models/ppo1-aim500-mi-only/0/default/d/317d4d2f6c3a4abcade3f68dc576f6a9 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6f9717e3c8590a82d1bf8d6e82ec27c9fe5a19c7d3fff459216145b2435012a8 +size 112873 diff --git a/models/ppo1-aim500-mi-only/0/default/manifest.ocdbt b/models/ppo1-aim500-mi-only/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..d2a9dc0594d99082f7b2197ab0e6fe57e483ee79 Binary files /dev/null and b/models/ppo1-aim500-mi-only/0/default/manifest.ocdbt differ diff --git a/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/6f764aa8da8d20f9790c34036f8ebf6c b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/6f764aa8da8d20f9790c34036f8ebf6c new file mode 100644 index 0000000000000000000000000000000000000000..86c21b5ebbfa1e11bff31b26bd22a673dd79c372 --- /dev/null +++ b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/6f764aa8da8d20f9790c34036f8ebf6c @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:798d4f8458ecc45374241563f68149da25d5e7ac96cf560b2194054a9643c6a9 +size 59101184 diff --git a/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/719b2023adb97e8e004b975110d4c980 b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/719b2023adb97e8e004b975110d4c980 new file mode 100644 index 0000000000000000000000000000000000000000..178dc3ab9b3e7ae8c662cd3e9dd6ff6da32fe539 --- /dev/null +++ b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/719b2023adb97e8e004b975110d4c980 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:792175f73e1473d2c4a53fc291ff7c03855f6950d9d8b553c078407842bd0890 +size 520146944 diff --git a/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/960c0e6c50a75b95799bb0d1f70cfe29 b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/960c0e6c50a75b95799bb0d1f70cfe29 new file mode 100644 index 0000000000000000000000000000000000000000..970e98d09ac07bb679ed747d8bea6c27ac253f9f Binary files /dev/null and b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/960c0e6c50a75b95799bb0d1f70cfe29 differ diff --git a/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/b8ab8a019fdf3cdfb451965c47d6b600 b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/b8ab8a019fdf3cdfb451965c47d6b600 new file mode 100644 index 0000000000000000000000000000000000000000..4b1a203a494f8d5cd524894f631d64fffb91355b Binary files /dev/null and b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/b8ab8a019fdf3cdfb451965c47d6b600 differ diff --git a/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/ecbe1ddeac541714c7d1c9627edaa576 b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/ecbe1ddeac541714c7d1c9627edaa576 new file mode 100644 index 0000000000000000000000000000000000000000..e6af2962f5caba1e866287b1d7cc9482791e9bef Binary files /dev/null and b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/d/ecbe1ddeac541714c7d1c9627edaa576 differ diff --git a/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..da1a382834ce74702479a2b3ff6d4434a6cfb1b4 Binary files /dev/null and b/models/ppo1-aim500-mi-only/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-aim500-mi-only/checkpoint_manifest.json b/models/ppo1-aim500-mi-only/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..1a011da2229f5880f0175e574b592f519ff0bbcd --- /dev/null +++ b/models/ppo1-aim500-mi-only/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 39000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "6c2ea3e2023c86a22c14aea1c7eca7209cfb840064710a4f5fcdf746e3715a1a", + "runtime_config_sha256": "6c2ea3e2023c86a22c14aea1c7eca7209cfb840064710a4f5fcdf746e3715a1a", + "source_commit": "a1729fb72b30ff66d9adf4d58b35a7ecc03d84bb", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "6c2ea3e2023c86a22c14aea1c7eca7209cfb840064710a4f5fcdf746e3715a1a", + "schema_version": 1, + "training_contract_sha256": "96fce81048151f1915df9c747b1a354e3e7af48435a6619d0ef8a11db8e73b1c" + }, + "sha256": "2bd33ceb37d80ba274368894e27411f2204ff7a9a535946fc1dda4bfe883a94d" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "19558e02016914d0431b33a5ce331ff05269801751c19d5843ce8827bea079e8", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "bc80bcf3a17cf1cc6660c83308e79796cfea225416fb2f3ebc3c6543bd34181d", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "6fc99613f77cbbb6eaa26bd935a6959a80bb745d81161eaf32df949150b1807f", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "179ef6e657eb372b7ebf3ced220bb37d0af0017ef7c9c1ff896aeaf501c36dd2", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "8a9a5c3a6a01f82d4b2f859ed02b6a2b1c07692d760c329eb1c06fc20887df55", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "08d446aa4be75198737475a68d8e0fdbfe0a1c0413788601eb68af1ceb09da8b", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "1ff5aa85b1856a870df808921abefd4faa3ef60f6c71c68d0d3fa3c6365141dc", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "0eb2a143ddb05b142c437271bcf92c3b0af8d22f3b8a9b3c6e699e402b9c14ed", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "114f1e55f8b01d52ed4dfe1e3f79f37c4cc10ce9e65916cdffa1fa5671d8a636", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "17e78114a5997b7406d823d14dd312402d487206b8d205746c160507d2c56955", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "d54bd63914ec23968a007754debb33c1d4ef0705af63eb924fc7a148dd1943fa", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "a8df5e5f7523266438b3c533b36a1c04d68039888b39d2f51ff68b817107fd20", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "1923668efaf936e1172828f44f33e740766d3d347ad9782f324a8628effa3ab2", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "14309fea6d2f0fe3f416c0fc6aab760dd80eb63808cb8f012166555bd1192aee", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "b29a72ada8d9f6f99f31cd05ab8cac99f0eea74e38a5d90cbf8a3f325c575b5b", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "d67e0c0d7963700fdfd5b9f7ab495a5510ab7485f62510f3de6f2e4dd5ab494d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "4ef0fe08a85f35e8d69fcc4ead7bb4425270f6a430f5a0e3f0743b8faf566eba", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "5127313108b33161f13fb285d7b5359b81dce93adcaeedd63bdccb4fb05aa995", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "7917dc970d6c7122ba5399007765c9d3ace60ed434db4f5d389d4aa1abb25410", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "e42112b4e717d588046cfc8ca5ad0ef92fd70da475217e77dc594a27f7442132", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "41bd4b95e3919726b095db2b5f5bd513815fb26d5a61d835b706c4eaab14df7f", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6948f6760cdd25c3da080e92234843f6847e4e53afcbd9da4613f467409f126b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "cbaeb01bed54f962b884802d02bf17383667aa46fd52335af8bd10275dd072ad", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "8f3ee1e98f3156d72d7f64c820cf56ab563afa48ac732997f016c922be2a43e2", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "2b1b470d3e4e35a8e7b9461e9c3250fce8d2c654a722b1447aab3067f1267118", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "1b11f3cba079f1bb8a66ee2204d9cef53991ba5f550926303c5b9e4750bf57ce", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "308c9e0c02d89a0efa75434b2be33307458341d87c2bce3b9b925059259c2ac6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "dd48847e4d8de9dda0cb52d0c9981604933832290fa14255d4634f33303b25ff", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "bc9ec9ee6d40411e636d75be62b98e93ed90becc43f0b01ee9f9fa0f010d12c5", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "85533ebd0546909bd50b0275753a1ea839e03b41f247efe7acd0130d3d61f536", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "c95385e0ffa7399ac52e83b0b3724770435de509c73ea15208c82a793ad2a80b", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "9e29e40957424c954a798013d56325cedc5d7d6a8ed884957c31711b7b959d6b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "eb0f2a7b84af7c6704dddf0b5e6871a3eacf85fbc4e6216e6106da1b3b896e04", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "fbd6f7ff5a3fc325b1f7f890e859f87fed654249022bb3e26b81970bc3d61a79", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d8fdddc9165cfe0b4e32e7b82a1cd2cdd413a5201c1d4948ef3a44ded4ae805d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "271d3379e84f8d868236576e2a72afa59d0f083be63bedb367344fce476f70fc", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "8b04c18eb853ece17fd89fc764b46676d046fb7312dc24eba0b832d016df6caa", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "46d98d0a2a03483cac8d7cfab7a529429279e65b098ea66d51fdbfeb86238625", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "0d6fc8e064d34bc9e8f939623192fa981f293de8ee705b3304f23511b5d4b866", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "590a789109822ac71c470638f7a6e2c104de194de92c4b6b46b158b4946776a0", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "03aad17cda6d8001c4f9ea13c48800f98ed2f9c2d1c644143123e2cda4b90898", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "bf9526828bc679ed23e3d6e8e2ee606017423f0d1ab9d056cd3da58348acc0f4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "d4f2dd5b598aa7985a10561c94725ffd2393090f045a8562d23cd0f89cbe776f", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "1c8dd5cc02d7be66b5b93d26ed3355ef72ef5ff3d2930254c0d12e99f7408127", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "82470e127b01dea6c4aff49c65a9458202054701b8bfa30a922e6a2cc274a218", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "6071556c5249d16dabbe629ab18a300444dfbbdd08c91eafd8000f283398c199", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "ca45c41b7c4594ed7522f9b6c990d9db7cb19e36a10eed98d1133e5ec44be243", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "fb442db0f3bea57039b3dc4c573b2baae140f07481d92a52b022477b81b19396", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "7cda1e5316f91a29dbb7cbd6446d9b4348d33ead163752bdff2edf4149896bb4", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "fab5d73cdf705f918d1f2ced4ac6d3becc0ae474e29a72330a24cee2f95b853b", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "f725092ef71eec77c6b9347202a951687aecc8d80512b7631097183a1c9aff2a", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "48d34426e73dc2a7b27a1275ff2e9e1a4940350fb176fba37cda6ee7e3ff32a8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "732be7fc5c4a1bfedce7063653fc82e35c41e1ae2d4dcdd060ffffefa1b61485", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "53b20c4d72105930af69c42a34dfe4aa5223b25054df90105213128ced2852d4", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "40d3d276195723f7f36f600e6e322883bb6f02442fd3b6dcc5dcc63ecf54239c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "0e404714e37b127e5f04b33c35aecb4b74d558fac4e95448214e338c6d36473f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "fb76f1c6853d3a1a7235c78ec945d516d50c5cebf3d613917a1cff5120488bb2", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "1862f468d48f8f623b3dba95a211c59069cdda9b1c1a09d038634ba0af77c159", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "821f302f49e838730344ed5479ee48bf9455a5f9c9c325820e23c9d345eac248", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "882f3d684f7365329c7a608e20a4edb5a1609c16fea7dd1d696352c22d334e8e", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "62b229b4025735db1f3fd2895e7a47730f37ad6d98f9c62d92d0e9adbd050f05", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d64e47ba73045e1789007328a175ff0bdc6d14afa258642cb3b2045091f57208", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "9ee84e3b0bed2ab74fbbd3838283a4f94b2c5d861f13b454770515216cda0fca", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "56616b5770f3796d3481cd0d4d93e59358447d3c3cefb4bf2c78c4033e673136", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "cd9e10657433224dd9dbee9aea8d508cc8d21c67b887347e5652b91ddf3edcc2", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e376de45a215771821728b425a5cfe87ac45696d67459f6dcaa1b3b69dbbf889", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2b5ceeac9294102d8f4ac8c86d110b18aebbe125a2847a18b65470859674c992", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "bcc52ec56b837b8c53dc9624f70f71b514177a229eb6612151f2e15a62c3704d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "0b0c665b908757c31830ff3507256e8480c7e36a34129ef958955733320f3f06", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "f6e2494447aa15f3a00ec9efbb8c38fac64529f2502d2180d9b6c78ae99bb40b", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "0dfe46de54c55443707bc22bd071d6b0c973d2ea9d12548c21986fde8bd6569e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "db2b7f26071d29f71cfa021e5e7bd27aa347990de551569f086e3fec7241ee41", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "c67dc5dac32ada6bf57aace97a2299b830174c943fc5ccad1057f2bdecd0888b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "f55c42155583bf9c41dc87ce609ba0ee08b6a08e5cbb42d990328b2bf1fbbdee", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "5c6b891886ad14348d613810d2dac4556ded7594b04c84d327892b1ba0481c69", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "98d1886062208d36b2e2aabfb30b3f9b9b44eb15a4c7b67c78760df685dfa1e4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "2db9212db8e4ca29b9950a5c9324163839fc543db6472d3ec3ee1cdd34c70f13", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "90c528d90cf15688eea7ab5183d5c0b0d3e1fab2d23595c2e71a17cf56c9f1e8", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "a27c492344b369c8cf6b435f208f08426379463ec8ce72cce84e188a67a046e7", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "58e0fcc2a4bc4fb47b458e032cbe0c99b57020c8e90c87e187ea8abf81cf1bf7", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "104f229b9d968c70b09d6eb8e31ebe884e3ea9c724d6340f1e9a9324759ee75d", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "201e9cb40e3723306ae602dd6afa4b40b68115b1bd76bb2b12d752cf5bc88f9e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "c9a03a3793a84254fc15c7f68a7f38b30f0fd910dae5b9ff22232d60a9cd7cd2", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "6c181d5ffb58084eb8ff933f753dd1cc3a7f0891258a14d6aed237649b148d4e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "0b279b51e404da296854555241ffa4529bb3af4f89641b35c1f6dbce026cd8e2", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "55d9311351590dabdded05f053518e14bd05b0f9f69991e50e93b8809bdd6cb9", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "f643a67bea920200058a5a042256cb01ed6503522605057969d12d7749758167", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8b48ca233342e1bd7afbbabaa1cceb39def95e563b4c80d493a41c9eaa4e13a5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c5acf0a4c36db067889df4b5425c7d0e7f5e16f8ff9c9331a53ebfd6d49831d4", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "1514f6af7ffc8d0a95fd83b338fc00755f3c1cc7c1ad0148bd54fcd4bfe8a7c3", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "19266d4ae92bcfbf6e52eed9b1fcd0d634ab58961bdb889deae3b9d988dae23a", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "56ddd912d77763ba022d3a29b55526c115c5013080ea7b01d1bfeecf75f0252d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "57681c300e55e53679f8487f089a9ff3e1bc200195119d75644359e3a08fdb3c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "5309f1046e7323f92077e00cb928619a84148a581df68bfee9212736e49da2a3", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "691dcc13f1b2a95bdc29cfa9e3ded0a91b36133109bb384be414dc694960ee91", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "febc06105a96cbb97ef457c9437239c0afc1924e1d113e8a9d5cfaed691c3718", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "aea7403aaf0e4ccdcca28364418058e761b187256ccd663402d9dd09b39cb5c1", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "06a078e3d981b8a4a8c857156014fa31b05c5ab2cf88098dabff86fa1a74cbf9", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "db7908ffc2c6db83c0f4db333f6d0b6ddb57adc1ddb2d2ce0c2e8fa9d1c01484", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "602491b46eabb8ee1f33515e03c0fea20892405aff2c8679ad8ea8154ebe2d09", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f80a422e80473bd2ee49fc1088a5b58135526311fcba276067f8b87984a8e420", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "20e06a1b6a483fae25b740425044d5e5d0e9ce4d47a05b72e3e52601c50f42cb", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e7172c79f1459a1c7ab1b2309a8cb1be149fd26904c008f4a26a4e6b1395764f", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "5802631dd836b42f23e07866f23293e02a904a3bea4df103922e10a3669a85ec", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "ccaade8b6308dc26c42a4844c4f7d678e587edad0852f9c0dd05228c021555c4", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f63af7dfb54e729ddac693b2f1504b4b785ed339627f2cf2dfa2cf4993860c0d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "bdfb5863b2c375806d2df0dd7c58dace7f0b3c18dab9fcb03aa7326b8f50c264", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "ddfbc58eebc65b165003e62473ccfa68e84fda64696101b420d38e597849e0a9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "07966639a1c08f95715c2c2dcc345bc4a98bd7055e015fe32074e122819f8b22", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ed69ecbb14fe1c9751f8728b013e83d313082d8c7feb0d77da4866828716a697", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d1c3c1cdec289d1b487f10cfc63ebae4dde7e046aca07fa4569ca757d5a6b26d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "d2894ce20ef030fb2158c5cbf5dfed3376e82df7883a0a00446855d21b431b91", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "21d76d5d0ed2170a3f3b74047673b62b81935edb67bd59946fe5c1ac4cf12499", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "aac43e744f48672ec004f3c5715de2f572a5dae5c3528c978901a49cb7b087dd", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "18d5dadcb6102f0ec66a6a8eccbb5009a16d06a211bd50763b75f558cf31ca38", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "5603c798cd5ccf7c9cd1a6dd8b32f7ba476c5349755216806d7f94b3079a1ed9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "746bf13986027884821b0235f965adb4837dd51c9f362b3707b4967224ded020", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "496aa69b664feedb07c17a82ed1dc0e0f238b74735308e3b1c89df92b25c7ea7", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "02e618a874344cbff6704cf645aa321a139b22c0f6fa68d0a35214a765ada819", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "d8deb27e9a9bc1890bb19e727afe9f6be4ab4990615217d24871e3688dcd59f7", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "2c2bbf99591cacbd5c4a0a12efc5dc83dbafd931e3c063b53b75b169b7c097b2", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f8c5d648063d2cda72e169921a44a58e7cd7ae583441eab0ba510b9b847180e9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "b47aa38873353ecb13e27704c459af95bc79a99aa85549c457f53c56e00201dc", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "391760b8bfb607802c87afc18fa05daa339edca07fb28e6f610faee0453c481d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "02c4661393603e136a7a97655535418b33f2df35e1dafa75a55b3ba303f0dbd9", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "37bef1286dc4f8222433857ad46018669d80ec7a0172f983b85dc5c2ddc0197d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "8d4a1086d60b6c000456bc8d6b62ff59a2487e25763171fa58b6df5ce17cb7da", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d822eb3002c64e5d5beb20cc0612a2a2ea733586a2dfecfc9487db10eafd386e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "7242462535586482120401f76cdccda9260bc2d6a2f37facabea6d248c5a3ad1", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "3f13908a686deadcc814310046c7fb466c2f2974cab59709b0cdcc05eeb61972", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "14663d8a21b2d6509e93e9491e8843e9f2843787001888b685761576a2b23ded", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "1a57476438d3c139a20b9d8e23f9b47e78305e1e3a67199396de21db4ffc9b07", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "853cd782540b4fddc0713c5019b864405787e79a43fc71df3f6f127cf8572f89", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "c7c5c7d71ac5595df0cee217caeb9e6699d69b133e199d01969875952e4882ee", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "b03f2ec3aa71257ad1ad4536045aeef375a5641c1b63634ae73d2ff665a6c5c3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "a29120adee6f546b87f0898994e05971849ed13eff9b9cedb1b948a839c308b2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "8739f7b5cf22af2a0511e890e3326912279a3fc0d46419a11181161cdd1dc88b", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "46439a684fa99d730a4fecf95d937de1094b0a8106f7233df5067c53b37c8aa5", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "79a88682f3bf3253a1ae690ec28a2affb505492208028d102a6b633b196b16c7", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "b0756884934b82b7f8b343a987513532f758f7a17c32b11e33957d7c3e2a88d4", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "0030b545ee3f8e05345e4d0041136b65fe96b94195d0105ac69bb613bb76cedd", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "7c42351b86d7f7958852314fe082e71fed8f99ca447694296218f394c1888725", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f3c5b03bf9c28e8a4ebcd1ad38ecd23985da3cfcb245c222b0a5fc877ec1f71b", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "10088a19e183b107f1bf86fbe968417e0c316269518b5188072f930eecbd737f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "78f4b04fda529673e6b81bc770134b82fd7955be4becee3636ef30f9bdb622ca", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "a322bf5b4239abcc68a600fe8a0913956e84b2f09fe85945484ae1955fc5cc5f", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "e1a2619e2d92a423660c58713d7d26c3bd56fa2fa5f4baae34c5128eb4efd617", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "2a5cfa6774012288a64c16d490c9982f2bf28e803c6c28642a08dc16ffa985d4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "cda2003e0a703e111587e193e75a68b65bbb53573eb4dbb11cb0ac86a855e1f5", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "a2671e9576ea6d3ac0b10f5641b1b91d5018632421386a689b1fe588fddd4d77", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "801b618c1636f9c863da2cd75d30d390049dba1dd66d1d8e6f992f24fdbc7aa7", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "cd344da33cbb6ac838ba738c21bc6b3ba2b3500ae71c8dfc30184c9370275410", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "281c68d32cf3ca282bab7dbc852b6cd47fecb8684529da748a6e293959bc6e21", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "27ba3a8e170ec3238c942f6c61cdd8f5f7b15b55cc015dcca6464f5431baba45", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d598f5ced8d7bd40f106041fca4d1385bac85cc390e7f7bc1951abd8d97e4d31", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "b1a89871c6fa97f6e65bd32f704a96808224582de9bf4541a260d44eb6c382c0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "f7c6e67c4d2f93ecf1cab745f5c686620f84e11f3c7a9fddd36456a68f04b4a6", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "65639a17c81375f50d54c65ae528d12d813dc176367524136531d0917b15c3bb", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "b5588208ca80e4af49bf4d3195e4e0b495cfaad85cabf36667282bfcbf094856", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "9907d2ea8ea5e30bf5f59973450032d7a7cd07d9ad5c4824ba3c0b580f083aab", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "4b3f16a46eb1577523aca79aaa9e9ecff92a5b2a16b7a4ab42a9c05b59557784", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "e00626b74eca8f0a9d7c5bc38df14d5e4c351270214fe3b46893afa866f18312", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "4e683d00b0ee0271d582f4a8defb88dcec0cc4068714d1499775894416837c80", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "97ed75ec09d585bb8b17440baaafdb027fe6d81bbde68233f343fe012b0325c2", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5ca5d107ad4fb0452b5e9b3b6dd2fef489d45031f67d657ba892245cde9063a8", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e916ac1dcb1be97e1baf77a132f6ae05576c42a71b53348efad5e1e9bbcb1686", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a444abc9d57512f8c6cccd864f4a0a4dddb5669ca160a1319aa73e877af875f7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1e1ac6ca47dd0ab05dc6a7d3148333228706a44a77d4ae94ee05c46c53cc48f8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "3421d8cba56da49fccc5b6ab1ed5c1fa15ebd12500edb2f460428804e42f94f3", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ed46393bd99f741d723f059ad2d985d024245205515bd06be7ffd49eadc2ee8f", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d060d6f96ec98b79db4ae59417842b345003f4b76a0f082b6eccb2ed1efa813c", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "1048221fb14cb7649ce99198e0e85aade5b4eb9a88b128e1a2157ce5d3e9629c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d4ee2ffc57b06df1c1fc6c2ade44f96527d390b222629f68b2f88a5a5f154f98", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "06b2c6b02f999adbd4e78de97e4ae823ee2b9d7698d8d62ad18ce3076c6677e4", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "5d977de116af20770314d3b4f30182ac3a764ca1878658607969a47f6f2e6485", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "d69e27877fe91c45757e0fd3f7b7d9cdfec3c94598e2aefa41b3b0b09ee1228f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "3986da844736ed3436f07d443fb3556385d83fc6ec1173e8ae5ec6786afd0e34", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "a0d15ba4ee5c1b31073bac65a57f77b4316fb18cb2c57983bd13e810564294cf", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "2b5667f72831995b0699fb645142d5917e24bffd43df19b9b8e93df777f34e6e", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "607b93437c556866995163b38a9807ce3c516cae6e2447951270edb3272e2e12", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "6c37ceb2c10e27540f74eff3fda662ae95377baced11d674a1740bcd317162cb", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "379f4d420575926f7b223f48e45d29e4b6814b348fe4958f3d09707ab3bb95cb", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "c12808ba12368b0ebc432a5faf3470d4a1f301f89c249f5547e19b537007c15e", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "6bbc9734cc3f8d0e26c07755a4b2634ad7ef0faca0af51089d128b900229a4ce", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "40ae71283e373564b0a7e2d58085475f3a2ad0c21c391c282b85f4d0ca807900", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "ac2b0fb51d5bd6fb97004a1d09241fc954d75cf34b45778fbac8dab6bae4fd95", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "d9cb805b79fad15341a3f68f9a117da13dba9c4601ecc8be58cf75c788336816", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "bc1f23f3bb149b72f3d79bd929b328a987a41d1a2ca08eb416ad898305796293", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "0ad764355ab0f00d37bd966f3bf79e16e5847a8607712d20d1f07aab6efe8092", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "1eb59898036683a7812cef04cad4d07cf243693e05b12f0c031244325fdefd68", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d81f5f9fdc903920af806a37e6b71d44296982e02f96bc45cfae9cce5d72e876", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7d577a319a7a0e073ac3810d39ef4d9559e5e4a8cc2f35ab2da2f706b4fb851d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e0508de7fad8d7b146c5d9e51ba1a3c3ec451dce994e68388d8cd9066d92eac4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "97e7fc0d01a7cf8fc9c345a96fb5a99206fb28e35507c07d10e384231e5d354c", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "f7df8edc03bb1690755f1c344e582bcd9bd8f4be182bf3bdd8f0850c500d9b71", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9ec23485072325b353ad0ae7297d4af9360ec3c317616d07278fc35c88afe5b6", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "ceaf218059ea3c3b8def75cc154b22258c638c59b4e4cc9ee626b416b18b733a", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "7a37662009f8db799d2776fdfd8cec118088267b22ad73d9da82939c3a845597", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "229c52c405f444cc94602df6c2ed537f135e01b5fee2a964550002a294f3b77c", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "b83505bc36983c11d5dcaeb6722534bbf8ba8e9958b096c1b7674f6d535812fe", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "4ccfea9df72c8b09fe50aac7428d87b23351dd56c1d4bf23ee2e213900af13db", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "765b84b7c882a83637726359f58c524a55b8771db50094038797a1576331ae1d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "046b67d053159c0ef5da85e2a75de6e3672ec5448c89a4c503ac7bad2467d51f", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "b5dbd55cf65dcbfe4e5a5c40a06d75dde50ca92dc3e0298d0c59de0065a206f0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "7f62769520067f4fd06db769d209ed102320cab9dff9f4be19b558677cac3176", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "0e2b58abbef7c383e1f7271fc1c01a0586b50a04cc9f8ae0f25b9ec96ce1df60", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2b0c1685ef983458151c6b405bfec39aabf10300dc1b0445c73dc06ad245b72a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "13384d13322ce33a6e6d958db631517aaf7e93ab7de5acc5eac98eb57e5f8402", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "f2213d752da565c11a5c099c21b7f9aacf8ec032200768ea96f9f37c8e7d0db7", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "c43a30ee2adddc8941797fa82799f4b7c9629ffef141844ab532e5998b9f831b", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "559d84613701e9cc7fceebed28125d469ceefadeb67c2f28a8ca96f7d42f66dd", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "9ddd4e204a9f4fdc2872a098665cae25968c5bf628674bae297f2c15d710fcd8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "cc2c2279a8c801ce37b4193e222917caa5f974ac482382c86a04972a20b0974b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "56b7e55be16fe039d47ab6b874746174e1d4357169692f5b838df3be52331784", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "52de5933414b1482327ac89e8b7431a9c2dc370215d8f90ea8c7fdbe27d9581a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "d434719cbe498812464173f5054e368cc9279dd8b515b25075fdc41014ef5ae9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "d967a536c3eb9d17ce0f9ba9fafdfc66d6c3d8cfce3b6d45214521f000100185", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "5056f156573593c1b801601d27e04ca330bad8486f924124aaf613d8ace50a4b", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "3c7e5500f2cd08aa4fd5f95912412849088cf1ac609873e85195b6b48b0996bf", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "205b8c3808daa69aabd2f6cfed379c4d734ec37f1c86e19fc55fcccc87749502", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "3cfc169d490144ae7cbf2c58fd7200a827942824389be9f1db06c27543080817", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "c0c0d6f4cf11a6acf98b28cf010175dc87a1e9df05e469c3c43525fa64cf594c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "8da57602f7a74ea79974d6c2b17213e2fd726abc8ddac6302848220afc3eb05c", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "28922bafaf1ad3a4b0a1c77e694b555037c668dfcb747c9deae62f9fff5e92f1", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "e7c28bd5072511040146c8339056d84bdf10774a46f06e514b9c2d1e7898ce90", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "8d6915d86c3bf5932131a415a907c074bada7c39fb33a2ebe2afdc630a9d4205", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "cfa008880df704b1750cf33a56668c276bbb661fce478a5a947a3d0fa730bff8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7603684666eb94cae49043af157ed6409422ee184cbc29d4b8830dce39bd7b49", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "2d889cb70d4464822d88939eda066620abb3d282310296419edb43eff2dd6040", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "dca7a4a60a95e4eb97492dc4b20c3befd417e76b3ffae55998d16e2f4946466f", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "546bd46e3572ff1d083b6e90e71a3649bd4f7e17e5f3367628c546820617991e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "2e82684be5fa7d34a97ec3d13bd4a77f407c0eb778568b3eae88054a2fbca3bc", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "0d2a4e7389e126f32e3adbdf1eb56042b1ef02dcf19330b660df5d835b5c6b89", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "73ca7f2f375da4f703e4ee2a252cc88dd1ed1c964ffdb3360c2183875667352f", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "09b473feab33c361c1c11959961079767fd1e0d9ed5fb68bfa08521a4b04ad86", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "aa368a6b054633e9d13461fb0dc056846fa94f1eeea262df8528228132b7312f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "f220647f5b6a34d116685ede6a3345da2be871ccd9d68f176f47e6f5eaba09b3", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "ef7de97bc9622f34a661301faf72c607a6fd95104814bd9137097d130eb01f7a", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "8595fc9fcc97d065d4e945ef0a68d2ce55258353160bf7e23c3a1bce8aef38a0", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "9e3fa0dbc3b8006270af79e0de193e9835d3c4d98447915d122ee8681e4ad32f", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "cbe1da4006c22fa14f423ec309df50bb636427f1dfa7d414d8fe730b990476e4", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "1a576e3b19f66f37b297ea10218e9fbd195f62c294dc8c95085b97079255b69f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e8ff543f4ce657e09d51b441a03b5f02a08d213baededbb9424d13e7433bd764", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "84dd1fbac1acde8a95ff9722e6391e0996d258714f9fbfabeffdec63392215dd", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "069324d13540ad161458e68dcc2dd737ff1b6ff6788bb3ee40e61eadf335c876", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "6b0ab7bab6c4242b2935d9cc6425d0d4761b9a61832e5d21ba0ce658dbd64877", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "c884ef3b7fd3104da51b39284b038c16da047fc87574d0dc93ea4a2bdf329d34", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e1d7f0e3122a8e654f31e1ee6029daae1d970b59d8dcf3a222c7c91ae8556686", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "483855ffe163dfe471911fbaf21a590e95ee965f1fa3582165dd0d7aeb7c4274", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e2e231eeb33ffdc1a0874ef5b76e3b365e2b387cb75949a3978ec3aac2d45673", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "0d928a0901d8688ec822c4debedd0708b361e650a723f67d313f3e060f161e06", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "39b2dc8931cb75c0eb224fa56e09cb0650de7151ecd5ecf123640d6e5f32e909", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "debf0ee523b54f227a949640882964d22647b34ba9e59ba45032b41bd244bd17", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "8912d68f8859e9830dcea9b5ecf516307c5abdbeaf6749636f31775fa6eb74e9", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "34f584e0ec56692b632f1462809c5fd453537e8e14c908758d0bfca5aa54a6fa", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "60429e98478a85131ece06b39ed623b394f4e0b031f4bf0dbc71186ffd739eb9", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "76ebb1f32fdbf711c7094e2d3a513feb134268fcbdd0f0bc4efe3c310c6a436e", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "97d58c357783468d43a5a6213d2eebb962c7f1d3aea9d24f67e883d272738155" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "561cbdb1ccf3e4e65d6c1d3a38e2f9072aef0257413ec30477f405a085c51b6f", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "a1729fb72b30ff66d9adf4d58b35a7ecc03d84bb", + "source_training_step": 39000, + "training_contract_sha256": "96fce81048151f1915df9c747b1a354e3e7af48435a6619d0ef8a11db8e73b1c" + } +} diff --git a/models/ppo1-aim500-novelty-only/0/_CHECKPOINT_METADATA b/models/ppo1-aim500-novelty-only/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..3880a4c337cafbda1a0f77357980adbc3c69b8a7 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786375422946965742, "commit_timestamp_nsecs": 1786375423602793710, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-aim500-novelty-only/0/default/_METADATA b/models/ppo1-aim500-novelty-only/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-aim500-novelty-only/0/default/_sharding b/models/ppo1-aim500-novelty-only/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-aim500-novelty-only/0/default/array_metadatas/process_0 b/models/ppo1-aim500-novelty-only/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-aim500-novelty-only/0/default/d/8a54313f46686b4f8386e4a65f7f301b b/models/ppo1-aim500-novelty-only/0/default/d/8a54313f46686b4f8386e4a65f7f301b new file mode 100644 index 0000000000000000000000000000000000000000..7d19c73d29bf2a8d228c13e1cf985f97fdbb1aa1 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/0/default/d/8a54313f46686b4f8386e4a65f7f301b @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:196545e3a7629f25f6cc74f51f2f4e6acb54fef851cd315146b5dd1b0222873e +size 112877 diff --git a/models/ppo1-aim500-novelty-only/0/default/manifest.ocdbt b/models/ppo1-aim500-novelty-only/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..49958955cc90793b20923c3e1401a9b7414bb65e Binary files /dev/null and b/models/ppo1-aim500-novelty-only/0/default/manifest.ocdbt differ diff --git a/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/271948bb29e4c6d53499572b3503318b b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/271948bb29e4c6d53499572b3503318b new file mode 100644 index 0000000000000000000000000000000000000000..767ad717cb51f98bc2b638256a68bb3cf536f600 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/271948bb29e4c6d53499572b3503318b @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c3c853c1521113f08dd93ee62f610d912cb715b6b1094b7b4164611201b1d110 +size 514379776 diff --git a/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/545a2abeb6a152cbaafcc1ad5aea98d2 b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/545a2abeb6a152cbaafcc1ad5aea98d2 new file mode 100644 index 0000000000000000000000000000000000000000..8881173ed2e9d1912ef32d31828b5667f47c4c9c Binary files /dev/null and b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/545a2abeb6a152cbaafcc1ad5aea98d2 differ diff --git a/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/583b331c85dcea84777c905344f3c9ce b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/583b331c85dcea84777c905344f3c9ce new file mode 100644 index 0000000000000000000000000000000000000000..03e726a451c61203597ea7cd1ac75a63da892fea Binary files /dev/null and b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/583b331c85dcea84777c905344f3c9ce differ diff --git a/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/8e59f3ea1b09fe58e0cfcb486371a084 b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/8e59f3ea1b09fe58e0cfcb486371a084 new file mode 100644 index 0000000000000000000000000000000000000000..2818e2aea3d04540e54e9e71400f57d137da22f1 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/8e59f3ea1b09fe58e0cfcb486371a084 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7654dd2b3c7f423957ebd95724aaa406bf1ca6c0b2b69a73d85028e99001f9e6 +size 61173760 diff --git a/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/d8205362561c89b0d3e19fa06f18d00e b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/d8205362561c89b0d3e19fa06f18d00e new file mode 100644 index 0000000000000000000000000000000000000000..c418e703944a613e3a01a64203ecd17ceb796b60 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/d8205362561c89b0d3e19fa06f18d00e @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ff362576b4387b4f6911da1bbb939177af64ece59b9f9c6259c1cb5e9a4cf446 +size 3827385 diff --git a/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/e1b1f6ccfb6fc53d6c9d307e38512ff2 b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/e1b1f6ccfb6fc53d6c9d307e38512ff2 new file mode 100644 index 0000000000000000000000000000000000000000..75e93b3315b28e71ce4bd7d20fba0ee3b771f579 Binary files /dev/null and b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/d/e1b1f6ccfb6fc53d6c9d307e38512ff2 differ diff --git a/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..0535df5f8ccb2ad95772c7306dfe375bb8c235f6 Binary files /dev/null and b/models/ppo1-aim500-novelty-only/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-aim500-novelty-only/checkpoint_manifest.json b/models/ppo1-aim500-novelty-only/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..18cd0e362510d07343fa3d7917869956627d91e5 --- /dev/null +++ b/models/ppo1-aim500-novelty-only/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 42000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "5f5c9610e51c646dc53e292c2ceb7713a7147d7fa7e89b0e8efa1a397889e1c4", + "runtime_config_sha256": "5f5c9610e51c646dc53e292c2ceb7713a7147d7fa7e89b0e8efa1a397889e1c4", + "source_commit": "51ec94a4dabbb366c54eeb1ddbc81b2f6344f4c0", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "5f5c9610e51c646dc53e292c2ceb7713a7147d7fa7e89b0e8efa1a397889e1c4", + "schema_version": 1, + "training_contract_sha256": "ec378f22cbdc7d8e91d9aac37b910ac5af4ef878c2495d2e8ab09e9deada4364" + }, + "sha256": "8a1c2f6f2c6c9f1ca450036818b9514c544a2ae2de9bde6231fcb4704dd4dbc5" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "d85ee6b24a5393e6a47035ced977e7c192827d673209d50aee0c282c9badeb3f", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "92009911111b739a9b63e78f1c62e1da40f11ee285c0154f7a0e2feed0cb0dc7", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "ee999169a95972c2eb486900973b1351f699874ca788bb6aa31338f8a2f07e64", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "4b33d031605c474d9e90ae034996e65c6abbdf3f7a9b4a024483019a79ec20f5", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "1704722aebc54fb62adf5e6a2e7b19fad8d7e194ee9051bbcf7c201fade23a39", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "89f60fd4724c5655d853d5408306391e7701920a718fd4324153aee9f6c691ef", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "a3e6b6a06900f4ec3bb49fa82a60833938570cb2b6ae3d412cfe1af7175a1ccf", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "0c3aff1463589839155246ee3b5e7867eab4453882fa50af333aec06adc7227c", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "aca31ea36370708aac62b0e1167016992618d788b29f5cfbc609d216ae95f83a", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "50d912c6e980701fdd2f49febf56b1dd286aa463e7253ad3f53629bd14574e92", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "05c8d211aac88c5699bb1ff8c1f481c3a7f15cbd7f6b67d9c322a6078141b5b1", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "c966cb3bc42ff1ac7f37694fb5f9bdc37439f884179e92f4c32dd986f61a3c56", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "b910fd39e435c5ad9b4137b42bbda16865fd02166a1960dcc1de0f1ddd631a90", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "46ce5b9dc144ef1c37efd225ebd7d47c3ebd6f86e73cb9869feba7b3abf35383", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "37096b382a6c3c08cdffbfaf4aa5404c61c964e8d4cdd6a7df464df74272cc63", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "f2f76107a6b88f754f7287acf942bf94de3ab9ef941f5474ed0e1dccacb0fedd", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "329385f01fa2095ba65769da798fb02eba43f95d2eae7d3ad653584c97d28fa7", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "326f16db16bda12e67662ab9fce9ddb465e625f8fb00de76a7c5024ae691f802", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "6d41e517cb0b7a86ef4ed74d0b334b59ae588b7e27b89e32d5502b7137d029f5", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "c376c29b4f91c76251f68e9cd21cef5dd5bf6a05c5e92f2a55ada5da8f547b4a", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "3d4a7e98446e3f530b9a97ea893b58bcbceec290182b7c176551458088155c99", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "5f767eab4e90caaf7deb518d84470d59c363e8bfcba784b23d3c8892eb483bab", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "8a4c38089510df8b44d214a6d20df15427cc6105b9246715a850a364b24170fa", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c21f368f0949fad5806c90c2129b6b8489eaf6b0b04ef79904bb973925037f7c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "18e6e0836017346f0858e361185024f44e1d6d65406b0894084be3b8f14df286", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "203230f13ec98b510f41f7421a708d02762fbcd45a5a1740d9efe9d3498e0516", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "1d9c73cf20bbb6a5a3290646ebd5c8600890d8c89187e1676f12dde6609c6d9f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "3738dbd859f6a87591956fbaa86c3b16633dad9e0fe1bf1e6129d5a3a6085227", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "52d1824c0e4a30132cca37433ecc3fd14035b18c959e9f1c05aeb870b755375a", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "bae2fe4211d089723c1707a49a01f43975fb457baaa76ebae1353392ad077222", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "059b09874f54c0bb6037b25487ec1daac6045113514df2c4732bdcb510cb57f1", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "df95fd132731370bd126a0499157d59b16d051f5a3eef283de8985571f625413", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "6904444f7d6a88b8203198aabe2e8bc5eb83ae0a1390ba4f39249fb5abe0d0f4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "3bc6590f2a0faae6d3d229d05f73bf8685fdc1fdabfcc397682c5d1c119123d6", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "91fdea0e38184aa270e9a355dda36652333b069d376b33a75ea11fc47c55f0f1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "dfea4158f239ebe5e93c627c5b06f4fcc4b6461a90874ec2c1048fdf41ec4895", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "9877a4a4dca73d902997c85caa7fc1617da0170291abe7331a068eff7c844f38", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "33edf0cfd21165378665a1e4672ab0ff4a33d035e84b8bbdc68815c1a9a644c4", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "b764e46a620aab8880e9beb55f34012e35c24471066b4f8984624f71b9965379", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "a3b5ed27dc2f9196a5eb6bb29f29a056c5c64040e5775231f5230b1ec849d8c7", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "42b661d3b8939f7a7d60d688c616aa7c463b66022461f3cb40ffac930e582a27", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d474e3e32b4cb25cae8516f6b98c8c621c59a194256f6b178095defa3dab4496", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f36d8722baf445c545278d8bdc9eb4579e25af9015e0bebdbdbe66ceb4ec7c58", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c7707d0c8469e5a7b8aa7392984ed835e14830f361a6649b5bf9fb6fab76cbb7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "a0f0fd578a6eab5d478e4b4e89f274e4a1158223265d8b43ac9afb5ab27abf1c", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "4c5a84dc86d350a0fc3d65f2e7c667d25ae3133b1fce242b1d8ea6928f35ca9b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "fe1508b887f263fa168f2f99340bcccc4f36b7b837aaa1d1cdb47071647bcb49", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "147302eef4a58bf1e9019c969efa2c4ebe3128afcc9359df1a56b42677bd97bc", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "855d2b8877bb9efaaae8049f82c129d97cc1df190828aee5d8ea98715a35ef3b", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "7497b980f20d80a5481fd70f5fe9ab3c8c8f2a6646a3bee57b77f84041be33a9", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "14876560277e160fcd32b4c34e7dfddab8455af5acaea075289e727cac69ff1d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "73f8f539afa382cecff30276e6434c084d2fd25ff4c82803fbc5581d5bc04352", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "5edbddc0e1bfc55c042f349e8bcdb1a4357819b5497679f4f823971d8b8d927a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "ae160608c4491537b0054d2ef2d9a9ebd0f70ff6e539f8749675e65609896a66", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "2233ee931f2e25410decc85563a0f6bbb802f01e22bf2e4c4e3f3d188b2062d3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "bfc7b66d9d4cd90950591515c97aa7348d4761e0542cad3eb34baba24a7ef421", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "38d169177f26edc679323e3a3e093289b66935fd875c1bb73a8e7f76ea922bd4", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "b9b7a6bb4f7206905e0fd57d4734db2263978ba748042b7f71371606ffd642e7", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "48d18b11c32ca1c99faa8089475f76c5168fe5386471e47a601dfde959341d07", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "1d858088fdbdbcf832c0dd403119574ac695f2ffab8e22956f45eb2f90882007", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "9f0125bb39018dece95e6fb82940be08b89d9c1747ba1c81d492f61b72ee7ae3", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "42ac8f673480633ba759b1858431b022e7dab5ea7cca565a7011835a644d1e51", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "7bb090620e1fa335f101597074956758fa50de8be508f27b059f57eea3505ba6", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "bb6ace9279ff32dbbafb810d670942ea3f689718df49a7b75854c79e39f2297e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "c82e667657c13a74ee94f71ce8573def88ebef5976e5411c9b17f1f14295dc42", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "ad9800bd9ed11ea985640dffb5ccd17d4ab32730651859ab7920c8f7b1b1debc", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "383fb7f27d33c9e696c929d8a7f8c837038e6427cf574f88a036bcb948814253", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "742fb2f7edc18c476106eb4535afcbba5f0cc91e9f373f691d5973cb86dc7589", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "9f3be7a713956329c60eaed04a9f88f565e73873fd9850fc57bcb72259008353", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "6063757ae8dc320f99c489b69e4bbf02657b895fca7faecc755bd48e227fedbc", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "a8823c6409c7c95293122a88e9711b2b7812829d649b97be6f41ba58c28d5e6a", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "d5f83c98da91efe69c8286d3ef41128404c88d2e5900c083aa4a021ef2cd6077", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "2bb0cf5fb2e58822fd100aa2164ceed04678581ae9957181cb9f7dfdd05de617", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "6b2279a7e8c03f51732a4f10cb9441ff099f56a5de44ba2fd23fe2a3fa59f57f", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "e7b369f7fb51a16a1bab0537a6c2f75d9d518c33bdafaa2c778791eb5daf507b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "97adbe10645657e1fa66135725fee8fa8dcab950fce8ed9c358c3af8c3f4bacf", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "dde7b16771a8404a13fa5e6d622e7dc771b0594c14e4a325de334572942d115c", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "8f2aa2b19f4c61bd6447583dfb7f6cd77f18b339013483aef8621735a3339d38", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "a139f2229f05229f241cc5b80f0ca5664e019ead303e71daac967a98cf7e8b11", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "d1e4b165340749bb53981258869952c901905e9a44970e53241e7a5b9f2f2368", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "3081fad97f89d6b4c2e8b4024e4ba90c00062e5e99cd13cb4d39651713c58ebd", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "230ec58058d63b7e6f839ea4a21070e7ea704a11d369673f605a0c373b0e6dfc", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "d030a784027d76649eca1fa9b7710ad680d47db5db04050fc232d3eb6159bd54", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "2e13008b0637d9789a6eb380045fad63dbcde7889ced5fd0146799be7d079685", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "b939a50a85b5c7e20d6ac9416f052f2325c861fbb6ebb9f4b123c0026d687bad", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "986241d8c1cd8659640fcd11ab7b6c4c6ef07826159d0acb3931e1514334ce1d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "02d8ce1163342abb7df2bdbaabc241f68041757bb0cba98bf977ce86d555ccf8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e8c55763c0d681d7147222257dc09e9cf187dbbdb6716943a8f4cb64305984aa", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "2979c5e0a602cbd7a2c71c66bf762cba95eaa5ef1d7e462c55369712518f1d14", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "2ab165a566cc2f20dbb3772347eb04c5d2a6cd14abb8ab1fde4ad445e2415df0", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "343b62ae9d9645f01f269d96a3d9e7c8f0b894333ee16234c416e93ed9802d1f", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0d6de0bfcea8b1d13c7af1ae705cffedb1d7779330568f13d1860af5273e937f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e9b193bb6215bca07984f5d209e929c6cef74b1b70116b611d1baa40c0fead2a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "7542aa85f6b3e8b22f3b531df7150cbf854717bcb4180f2f8a22304c6d98276e", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "a70b98681328f7dfe8898fee4e23054769b8577853ba90eab94e78dac4ce6851", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "f7363d4455e802f7d4826165ec4e057552e2fd50be18ca7ca0fc6c5b870ee798", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "940ef6412dcbe0a65d7e54c91fe4dde97dc344027791e7eccde9c63a4469d0f6", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "52cb20b2b4d99394dcdfaef734f8367ed4412313d2e678bb93c858088fcaaa72", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "ae34f7d23b7ba5148050f7a1adfffec6161ddc9c9c2e359a13c0c2ce8d1182bc", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "251828b9a2d76cd1ef758ec9b3c20db8dd5af304f39e95bdbe80b1d9b98e2956", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "14af4dd15a71ca8ff62b63a9905b6b3ca87de4da052e05a8acf18b721dcf1476", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "77840442fdb3a83a5946e894fd1c4948da2653be0748a0529b8f7d86ef60a1e1", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "d4f86d501133dc16ed0c5044d1fd4f21eb05c01a373ca19afb74a53f0884fac5", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "8d3a88c25a28bd742852631b0eda2e6b6611fe69a330ccb7f9d7b012332916e7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "698bd366c72d3aac3e6dfc7dcd22d915aba3df2dc6079c76f9b4f43a87a33c89", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f05ab13b0c714a455dcd97c0ecd218dcb7bd7a15fef80a9122fdec45de3593d7", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "e8901d252808c6ed49d83244e4a7843a3fff889a2cc4e25a5d1e37e31ce637c1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "c9412807e4562721faad1846b9d71390b40ae5270c0832c49cc0e05a21146bdd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "b56888f94f5bc58f3bed4ee23832384196962337894337154e2af3b7dc73b71a", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "8bfc3ec8e2581bea17da79cdd421d7a34d75ed9ac8f9ad5643e7c7172d22e706", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "fa3a110b6a58ff5b08b398f0fd5244e7d8a41c573af2bf85f4c3dda5f5d1ea5b", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "292001797769be12cf7cd5ec4f83dee4506908c12dbfe684581886627dca289d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "660e9cf9ee47ba6be1b25a7a75a499ca38badcf6c16a6ebf279ae688683c8c68", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "25fb52735567d527f7132e916eeb733443ebca3954f9cb7f23a2435b07bf9385", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "69aedc82c2ebded88b303a026062dcb2dd81dc1cdf0e3395634ae8075a2f276c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "284d43a460b3ff1b688f292cecb54996e528db899be47f2afa8cb8d426a2c25e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "9c1c88c10276e08f9b4a989281576de4d30c37cdd10fd336a52b34041cf7695a", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "12b20072ed6875761c64b40b31a7ea4394a5c409ba41925ffef9af029c3dd6d3", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "5094970b9dffa39433b5a59ad46f55498c774e8c825c4ec6658b32f27c83cdb4", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "392d94fc7860495e067660047bdc35e94281fb5d6e936e9391cd034cbd03e4c0", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ce837fe4498eb9eeb4d6d2a150c5c9a9ac8e9f773329880c985e4723d6c49669", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "a05426fc822b38a75fcb0a5d6f173bb8be279e2922794675ceb22c1cb5d0c14e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "af9b45c11e8cffd91a7885389eeb85619d66075fce23403c387de9ffe3b2d210", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "f67b6fe78de6d59a2a607efce5deb1474d144164c9eab8814714821fad15bff7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "41faa183800918df3f545586c454149a5b864623e01a83412c3eb5477c1daa18", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "19185543c2611759bbb535d71f44f34e568c00f5fe6b33b159332377dcba70a0", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "7a041972169e9e0b14fe0591b999e52dc34f61070d833ff41653315dd1256af1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "852e0007738f940716dfda7abb2d18ec733dae726af2801e08ec8c17872ecf8d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "afe45a03916d69f120fb58936f7f1cb301d54452521aae740ef5c2a100800ea9", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "acbc253189dfc1863f4f1890eb76cc9a5c0b506b4dd1a11aa6ab5deca76efcf4", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d1dfdb72063a0873e0baff9bf860dbbfc5b333d8ef306c9337a5d2e0088b1fd3", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "c62e919ffe1e31cfc4eb3cbee7a8442aab4f45399684d2773cb3061fc48bacd1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "00130f0a2c94214a39662c6c68d0c5f6b61e0c3912742ecb84db08b2ab21a910", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "7b8b073c10832ab6894655b699190e413d51f6f64335aa3c79da95f060cd7f3a", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "2af5a16e81fb12963198745415e664a175ed6e9ce7b99b2a996886683cd28682", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9b7139829994fec011e3f39dff493ded9268db6a72be18ff674612bfdcbe9a9b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "e39c534e1dc7aa5f631c4bf97a5eb834f2c59a03d68f716562d9d22b0b5d8acf", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "05a96d50f30aa9a73d8f59bed4d4b17e24f8f4d1da09f2a0f4d39e4c3064f06d", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "03f4ca6eb49910857169f46aeb7138cf89fe8109ced25c551ed2c5a98996b803", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "26d573ecee92ad3f3735908a928a6c197dca39e4e530cc99e0adeb924bdda164", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "c53aeb622a19fd4bfbb9c0a0cd4ddb3da953bd7a65b602f72012c2670832499d", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8579cc81c5b2be02921846b7ed4fc16e3fb540ccfb343e4a7c75fc6de34394ea", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "3d5b992c3bfb562b8b64376e3b78203be838dcc414d0f94c369a1753e845058d", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "13116a26b66f827bcf36b62c5ad00ce3720566aa2c1caf5337a28ee9e21881de", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "951d2a6aec8db1e5d7802ca3c95bdc4c1bb0d0c34409d73651c42775bc2d0e90", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "935f6c3f215bc009e78f88371aaf401f2581e69fcd99db46f13a694f249c8e56", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6059a903c589baac009eadd3b3d91223cb8c7f6deaabb064e327d414c9b8ecb6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "54a6654ba7e6084f9ce9ac22844ce5262e08febcb9ca9a1b12321df5a668fd00", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "835626bb8f2347885cbb8d82141c0b7049079df6eec06daf175dfaa3539b9355", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "14e537f5f356073d1efe742f591ae84176735b80d2de0c60078921bdf059fb1e", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "6b9a31f27429b77328d928f37d0cf2865d2a3129b0b445fca67a9c06fd8e3942", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "bf59f5a9b9663db4f65b3fa8b569a6df9184898475933c0b83fc199b54f54ca8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "9b9576f17412250ecfe43021a9b857ae092ccb92e1c15fa9bd61415aeb87764c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "6f951f3f602648be3ddb089a380180629b10282929b9c4ad5aefed60d4c5b829", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "2da20d1426d769a00f0b195fa0c1b71c99de19a971543beaa49e6d37fe50bbd8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "a945720fe3621e7f92c593522b0edb795107085f1fa0eeae6d392a4a2af75209", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "867f9f3e378379ab196006c3929134d59fb1add9f16e5a16b3bd25022ab18c11", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "822a1f7c1596be49222b0812200647c22ab45ec51bcecebd5bbc8f891a16e933", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "a1e91469eebd433f3e97cd708acebe6287afdf08f04e171b07da00573f5a4ada", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "fb1f6f2820073bc561ca53680521525e4b8e3914ab82187a11b3168af661ca50", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "b8e07d999f1301fa618a5888803ee9b520c6f426bc2c549245c98437d6aecff9", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "73941ede29d1ca3cf7416cdf8bb2559f63c1369f2fd59a85974c15b2e0a80b81", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "2a729489cd8325c9237baf4a2167564d78e812bc2dd8adc03edf91826ba707eb", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d57992b638cb10e59bd667949d46986fa06667d19438ecb640b7c70b4deb878a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "45189571b55b776dd2cee34346f311c71986f40972edd12a31878336b8679ef0", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "ce48bfe8544ad5e95b49ec5f6beef61ab228c823ab24373cc64cea38678712a9", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b9831168592d5464def65ca348bf53c9132f8a1a391b70807daa7d38a075a0f8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "cda1d6caf98d86cd297fe827d8c00add7ff91fcc81e34a565af615d0ace86e1e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "6ad6b3d0693c684f19ec663fc7d8de3a7893706a7cbd1507a555387e679c3a02", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "691aa78e3e57131e19c744592af98cb882884e223b3a73bb7e4a1fbe2f586be0", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "59a88eea89ac1471f754664df571cc211f3e3453d65a7e6dbb76a7c23958a8ec", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "fb63efd7265ef1f2dddadf62d64b9d9a6b357aa3d6829e65467686ae8eff2cdc", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "377d5e80237172617b4dfc468fb0444d5c0deb763ac617fe0cf00e6528a3640b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "059cf020e7bfa6f15a05ab67cae0f06a999c59ba4b0024db868b7f977d96997b", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "b8f15d9722612bff82badb89c07cd702a21bb99e59e8fc0293dc06115651e94a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "153a56731312e9123f40842b73f3a8ad6193c7320c9f20bb35dafe815b060ad2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "d85795afa8a4d10c3c583fc052d27dd262338fbf4d8aeda650d47590e1bdaded", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "f30b4651e76071b7845e8fea4340a65a06a0fc6b59d547565ae084a9b24ef101", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "d30e583e0d8ebde88a3931403053201ed07800bb352fa0bd1d109b456c1fd2f6", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "8d98cf8bd0c032961e00468d18ce54e6acbec66ddafc35a34820ba64a857a4ff", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "466424b2acaad1f1f00bc0ceef021d186d8b62f4fff151ae7c94380e63f4a45d", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "9b25b2b27178297d2a0a1326d40236e127f79f3762947d083434c004adf25174", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "44b30f97fcb1dd8585ffaff57c52c4f9ff93f2b0409877317d2c01cb118b9331", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c4877178865802cd01094aa5e05b5213af72082da8a7977744d1cd7db3068de6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "be590c7549aacb78c8504bd8c7300ac3e40598029d628c7cb8078bbdbea1d3a4", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "b43085200367470fd9ab3a71a3932b4904e3b48ac8f7ff2394a62a2e7ac0f031", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "df00c505ed36c45206b0809e24d90910ba41925982a496b26c5221fb9a0872a9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b958060a7e19a4317fd7d650c51a123eb793e62cf76314e5556fa5688ddcad4c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "6f5674cf7cfc3fb650b3c63b3258f5a990e412cb4a07ff151fc2c6bd5edc2a5f", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ce806bfb5e20041af4454e5b54598e288711afaa220490c55c1551acea039973", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "0cc64c3ea08bf744d7a88364f0f0c6c870c7e7cd7d2dcf751997e8ad3f9cca6e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "2439c1f86dcd18627d5ea98d3c4b54355b580de64ff6f3769207b28f6b3dcfa4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e63851a44b4a053e8b50cf454a8cfa5691da4cfecfdaf575752333f7f8215187", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "ca55c958a50a197cee628431b07fb7a2b16a0813cee859e1abaf36d187c21d75", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c075c259baaffeb922ea884379e5f6a7defca2c7d9c4a2aa66efbe7901b4e351", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "29a2fa65447362e747157b0197572c0e44d429abed04a28a4f828cf2ec09f088", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "626f7d68374041fcea8ccaac7812686046938d6845cacfdae2a4a7d235eefc71", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "23037d2c7fa2d62c2e8d7a7227e8a3e8d8e5b852ba18f8a5d24e3bf5430e9a73", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "15ac5bf1de0783810ffaf61b2285ff77a1762d266d9ca4874132bd74b7a56a0b", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "5cc75fa0dcb66d4800c12b613e9f18f3d43992256be4ceb19226ffd681896f63", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ac39172efd62c4b19b8140bb6d273a698e6b1e79264ff858505bf4ecacae179b", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "fb4f3a9ed9349ce6f1ff9492aeb136dd16909232637cbb7c9b0c6f0aeeb0753f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e5c0c2d62454128c0a910a067c4f3b7d72eb42a167f2d431b99dde9aaa592746", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "8f7e06b3e89e8dd8e62c8ac1cb3681ea24ed5625436b52edd4b0af619061f54c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "d24a2edf3a0c7946ef5492c7b5e414eb51dcb25d7696ce06aa0b7fff265a91bb", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "6a914cd4d96e1941dd1226d843a095cccc0118e51101049710e254d7e2d74b6e", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "3001d876b3eed3c895bc11ec720a70c134e4725499ebf06d0687d77e38c26b27", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "34c221cc6764934200e09b58149ab64f99a58e72a66c09409a120962db9cdf16", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d91d3e6fc26499ab376795e73e94ced9a77279f4104bd84a5f6c2db9e999bef7", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "7cd08af673a17c3ddb3cf4bf16fde009749231abbfbdd710ad43328893efb376", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "52d2f9d7b6e5ef7ecd37dc3e0d6f1cd69b6347b72cdb2ec79b192480bc56ed9d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "f7a7f5171742184fcb3e2eff93d771cdbfba7a75d3ead7106ec3197b8d3894be", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "6c7f445f264645f85c394ea7d2cbe53abd02ca9a16a7e143b5474558077230e7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "8bbff1588eefd1a7655afce53620cb1b5fe79dd3ce2d163f7458f39eecc1c76d", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "9e07aa25ac2436d0c25c2bbcf87f30011a47595b3c409ee03b849d5e862a1a0f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "2f966be01652f12cc954b4db034f0db01540a5816530583ab5372f2fff6e9018", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "fc7c54b8f31619c3ad24d1ae4b6d755abb4a37ef2bf4bd5d150838f96ce44baf", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "1823b478e7da7aa29bedeb0fb0e6406a1b2d2fbf3aadd67292543a247629d6f7", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "3a523e9616e06666e6b727495dedb3976a83a06fce157e930e939111d28cc9cc", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "2bc33fc02148c3e3c7065b6d9728bce8944c0f8df68cc78590d0c53d5e9eb221", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "aa44642ea8b4756be1fb082233792c1b8cfaa491d838926f70708f6ca65c4760", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "27b021cf055b66afec177455f51523d21302e3116d43d66cfa934d755b741d2c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "459590241509c40c3690f6aca4d691dcd082035a74751a008c8c7072ee566e81", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c69ad5257c51c977d8f1eaab4005c113503dcd5e949244f7dc5b67ff35e582f0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "173a894a01f423da4ea6c330060e1157881d42233ec44774a173968d3d6d1763", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "77c1b69f782db9910318be86831f6ca04ebffcebec44be3a204da8bf132d064c", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6166ef648e4888c057e43ba79a41b949de377d4777f1a188262bd01b48a08553", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "5e33e8853228e41f3dc6a434ed36df23909f738de2474eb848bf10d090cc755d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "1bbc19b12e157b8a1577c5f696a18a07305081007237a7874d4ff8bdabbc786a", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "09575784ec542403fc8c486bb31c398cc3cbe6e7da8708e35701650618fc9681", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "e06b58973ef5e08b264901a31d6cd3de5a9c01f7067122e5acd60fb91658d93f", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "714112f56ac15df75594a9ca7c55ac0bd4b880cdba61f9d3ac0f6c491e5f10cd", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "67767c2a02762d5917aef512cac57c6388b82f650d7c54fb38016ddad9e9074f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "69c7ad939f21e864780953db6b4139d0f6bb74a3fc7bb77730b2709be5ff0d1b", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "e86c0b1f05da1bff5706aa88102a57f6571fda9c261b21c3e6d8b1685af7fc8c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "0869bf0df93eb5728efb3696e5537f4564d3de4a1b5f0a6e0cc46c5519957951", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "35e1145164d72d17695f248a8f16941f158c32e86f2979ede8e1f60624417d71", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "587314ba6d6b5f75b1fe9a6105d16a35cdb557a4d1eca4f798348c8f4b55f679", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "ea64e451850616535b475a0f18bdf85ef1dd45d2955fbd22eaa7b9f1154744be", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "37b5c5891aa8933f1f3c942c76395a525fb0b69240b69f7514c80cb5f9891be7", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "647b8a17a2b76dbd9165ebe7a36449383115e2d8958edd8a2310c7c790aae660", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8960b1f0b5dd3b06ae2e72e0eb71e388ea8e115c2032648fb3aa05d844fd9d0d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "279dc6f91667c011d962088b4730b79b5a88d0c47ec89fd413a7f9140d7c1948", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "e438343ef73a93752541e052a90b66dcd2190431ea199359da06055fb062faf7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "084ded5bf3c8264f38e8a58e9991be47c4d468945c2e7c21fc2c2a23b8a7c9a4", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f278a81820c27db311cd19b70bab66dfc5fb4dd21255047b249e7b2ea4ed947c", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b9bf5ecdf55a60aef6f9c7b256a0d9cb083946f649cfe63f787d38f3059b030d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "01fa40bb67e5cf3ccff5f8b1c008e0377549c47c8831b24c0099d2e8bb149ee2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "f73966da3971ee181f0957ef4f8033acc814888404bfa5c94cbf944ebe76cc47", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "69fde52a2681d678edda6318c104b429c4a02e56a97616b33b52e95a45cf7788", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "266531b3da2c9d698a4c9d40f50efe5ebd00e8a9aed27dad070b08620169ad37", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0c464c29764e06e244d957f085352d5dbbdfb3a207a1f8a8308af3e0648190e7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "225c79a99463eeb18111663907ac90ae9867253f93536fb3cf5271674c2a654c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "70983a2fcdf204351e5cb7339d3b9b8e964861de308fa0c1e7bafa478353bb44", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "33b7c55b95804d8d5b8d9b2f487af8bd8fe32546b1fe1a12c3507907928fe860", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "338d7cd6f2741d0bf0e42f2c4a1fc27a4693c51da04ced8e2020233fbae30ffa", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "56150f1455f7aec4082ec7938f2f8161a3aa07320b1285b273113871face7baf", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "a5f0cd0062bb7f90d774ca99d8c975366a022e93a3c99a9225f2d5351456e8b6" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "5975e49b39122fec9983acb172dda693a902ddf574750920b74e56ef8e67f4dc", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "51ec94a4dabbb366c54eeb1ddbc81b2f6344f4c0", + "source_training_step": 42000, + "training_contract_sha256": "ec378f22cbdc7d8e91d9aac37b910ac5af4ef878c2495d2e8ab09e9deada4364" + } +} diff --git a/models/ppo1-dads500/0/_CHECKPOINT_METADATA b/models/ppo1-dads500/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..f3390472f06a64decc8d8575fcb774cac771a25e --- /dev/null +++ b/models/ppo1-dads500/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786263792791668244, "commit_timestamp_nsecs": 1786263793448753255, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-dads500/0/default/_METADATA b/models/ppo1-dads500/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-dads500/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-dads500/0/default/_sharding b/models/ppo1-dads500/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-dads500/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-dads500/0/default/array_metadatas/process_0 b/models/ppo1-dads500/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-dads500/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-dads500/0/default/d/d311620e8f63bc8eadb425e79cf903d1 b/models/ppo1-dads500/0/default/d/d311620e8f63bc8eadb425e79cf903d1 new file mode 100644 index 0000000000000000000000000000000000000000..d3809769bff2cc0f22e353d84bade61252863a00 --- /dev/null +++ b/models/ppo1-dads500/0/default/d/d311620e8f63bc8eadb425e79cf903d1 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:77044f7a39188f39af0e5ffc30d69633dc03809cde07524e7091ca6fe4a69470 +size 112879 diff --git a/models/ppo1-dads500/0/default/manifest.ocdbt b/models/ppo1-dads500/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..09622030298175c2d35a4b48ea27e9db1d8eb6d2 Binary files /dev/null and b/models/ppo1-dads500/0/default/manifest.ocdbt differ diff --git a/models/ppo1-dads500/0/default/ocdbt.process_0/d/3ea8182f4a879390c95bbc7a03ea0c38 b/models/ppo1-dads500/0/default/ocdbt.process_0/d/3ea8182f4a879390c95bbc7a03ea0c38 new file mode 100644 index 0000000000000000000000000000000000000000..970e98d09ac07bb679ed747d8bea6c27ac253f9f Binary files /dev/null and b/models/ppo1-dads500/0/default/ocdbt.process_0/d/3ea8182f4a879390c95bbc7a03ea0c38 differ diff --git a/models/ppo1-dads500/0/default/ocdbt.process_0/d/5052985b5acf0be41884125bcbb1d120 b/models/ppo1-dads500/0/default/ocdbt.process_0/d/5052985b5acf0be41884125bcbb1d120 new file mode 100644 index 0000000000000000000000000000000000000000..1ef4f09e7d501fcff3dcd15deeaae86905449ef1 --- /dev/null +++ b/models/ppo1-dads500/0/default/ocdbt.process_0/d/5052985b5acf0be41884125bcbb1d120 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2b4e61a9f890ed1d9eaadc4937f9907fa9aa84a555a3bf5f6fe64ca7376578f0 +size 359825408 diff --git a/models/ppo1-dads500/0/default/ocdbt.process_0/d/c4a13bf5b23b4f3a8bba02e147f22930 b/models/ppo1-dads500/0/default/ocdbt.process_0/d/c4a13bf5b23b4f3a8bba02e147f22930 new file mode 100644 index 0000000000000000000000000000000000000000..35e9229c99d852787efd5006fade2b2dea8bead4 Binary files /dev/null and b/models/ppo1-dads500/0/default/ocdbt.process_0/d/c4a13bf5b23b4f3a8bba02e147f22930 differ diff --git a/models/ppo1-dads500/0/default/ocdbt.process_0/d/dfd4b6d00ae3a8b91abf37c969f53f9f b/models/ppo1-dads500/0/default/ocdbt.process_0/d/dfd4b6d00ae3a8b91abf37c969f53f9f new file mode 100644 index 0000000000000000000000000000000000000000..4b1a203a494f8d5cd524894f631d64fffb91355b Binary files /dev/null and b/models/ppo1-dads500/0/default/ocdbt.process_0/d/dfd4b6d00ae3a8b91abf37c969f53f9f differ diff --git a/models/ppo1-dads500/0/default/ocdbt.process_0/d/ee22700a49c9fa9af6bd1b82129863a9 b/models/ppo1-dads500/0/default/ocdbt.process_0/d/ee22700a49c9fa9af6bd1b82129863a9 new file mode 100644 index 0000000000000000000000000000000000000000..0baa44fe659d4a2ae25550298dcf0508d7a9aa6c --- /dev/null +++ b/models/ppo1-dads500/0/default/ocdbt.process_0/d/ee22700a49c9fa9af6bd1b82129863a9 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e009ddcf61627bdfc7dc856c5e36f77d73e9822ecfdabcd1fde3a3da167b23a1 +size 4382720 diff --git a/models/ppo1-dads500/0/default/ocdbt.process_0/d/f289a04c53c6079360b246d05472938f b/models/ppo1-dads500/0/default/ocdbt.process_0/d/f289a04c53c6079360b246d05472938f new file mode 100644 index 0000000000000000000000000000000000000000..152e3c55c142a29aec9bcf78a1d31b1cb403edbe --- /dev/null +++ b/models/ppo1-dads500/0/default/ocdbt.process_0/d/f289a04c53c6079360b246d05472938f @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:86df92ce262963b74fabee9fc20b4608722facedabcac49d3474937f68f5f6b3 +size 215097344 diff --git a/models/ppo1-dads500/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-dads500/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..1b3a016f6a8d91a231f99cc199a5d005f1fcdd7a Binary files /dev/null and b/models/ppo1-dads500/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-dads500/checkpoint_manifest.json b/models/ppo1-dads500/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..65ad994d772bc215bf44ac2f111fd51d5e4ff196 --- /dev/null +++ b/models/ppo1-dads500/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 29000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "7b556e8140e7762d1200ccf924dd1bbb6e97567d20ffe098d7adf579fcd14f03", + "runtime_config_sha256": "7b556e8140e7762d1200ccf924dd1bbb6e97567d20ffe098d7adf579fcd14f03", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "7b556e8140e7762d1200ccf924dd1bbb6e97567d20ffe098d7adf579fcd14f03", + "schema_version": 1, + "training_contract_sha256": "c32f279d672acfa2d0ce7c6cad11410b145c081a7716c1365909274e0c11b0a2" + }, + "sha256": "3e5a9551994be7e5a4b27d5472739d1d12c8a1a712dbd8538afd93c7d866aa8e" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "dd409cf462bdadda9496133888de51f1702a99c621b03f328e71393ce5528c99", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "f9cc5b3b6b1a0fe03cd932c47887b7f4af5f36816858873de53210bb3f6f911b", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "12f010ed023cb0656f7cdef02da93865689b0c2f31a9d9101b28ed6a3c8c60a3", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "2e815b5358a1f33616e2f76ae4e4706e220838fd73dd10b6f156ede0276bb6ba", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "21fbf3461a0f59d20997172a3f1ffc08cd97e16a989679692a8d39d72f3dde0d", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "7b43f34685a8323023cd22ec93ff8c5e2f523de1ce7bb615f6df850fd968b043", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "b9420383055b96dd1dc773de8c89c3d307b94586b4fba981d4be4462b06c767b", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "976519e4b675035b9f905c235c7e282899c839ac1b74a8511f2d9116977d6b8c", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "9ac83d9e36454b52b91beba3c29041a1f0121448604456e1a53433b8e3078e09", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "cf62b4d5ec77cbc37a2f99993b6b02985634b8bbe20d293227948e911e172622", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "23c7751e6fe364b7a6b11b5f6ff62df2b51c5af3ed2b46343fda777300f2747a", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "39337a94eefae296e3aad58958d1b794765dea9848cfa3996f31e1c45d5217e8", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "09f4e3e11a7ddf359c033620da7aa711105f462e70137a412bab23e396fa55d4", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "268de4cb29cb1dd6c5630f20e96ab9c76fc0e691b2c2857b77894066bf15f04e", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "fd116382d03b7cc5e0e048832a9097b2e98aeb6eaba36ed068d12b09a901707b", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "3239a9737a0957d9ce9e01b021532331711a3b457a4befc10bbf67d7fcb77d44", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "25b4bf3acdee4b0c2b40f383377fe3e73e002917ec7b2309d8a22a0c81b4db13", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "a941bd7ef711a70d3484cedf7b2ca5157e35ea78960020ad9887fdafd48cb38f", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "3701554f18a1331649bcd50c0ba1399efc507436e0d8bd2afa0e5764e81835f0", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "e1e1d862c82656a833e24aad9648bdf1842a0d69e3b7edf525de23a4e75bcd54", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "3289ae3bd09b325b32a40b09148b8300fc3ff70664160a65378b1f6aa9fe1c06", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "800c7512f9c479c615cdae40b8d25a893b186a3e08c816835fc579fc4ec5fd20", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "fa7309c4d75f4a34f506e9cfaa04bbebd1e191baf6cc6d72b48e693c66027dae", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "b10694b38dd315853469b107a10abd52e056fccac1214a9fc509964bf3ec1ea0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "972d83c17eaec52de46d380e561f7fbd3c7ae1755310112a239f7944bec314dc", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "cdfe506cffa875a398ee6955302263c5b6ce0ae90a685aab36199fb993121604", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b61ec53528d44eeb0e75a74310705c2966b242a33eaab23181cf185e3c7ae644", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "2d6cac2509eca178ddc3bf60a37418a76f14dff6da035511433fb4d1ecc99c40", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "478110af5324ba4444ecdf27187c622a0fac22c8587cbe96550629c3d990c746", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "4efa51bf215bf312579cc8472db59729d07fd2f5d2625589896669cf48f6e0a8", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "c196cdee6f36c20fe1e940dca627c95edcb33d496fbbaac4397d845fe39d6362", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "4930a4d6b4477d759d2dc9323b2409cc7f54a40af091ad2439b7980699938030", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d8271c45887ecff8c0b39b0f776e3615296f1c17e69d29af519e1404e809f0cd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "59afa34cf09be1416ca66be2f8ad24e58bb77311bc90b8cfffd88979cef642c8", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "af8913800f6293b99a22c8eff717c883bab5426bf7717eb63cccda1e1386e747", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "1190763b51ff440aa41ebb64abb5f9f26f116813350dd014f7884631112d2c82", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "3104ee252bb49614f92c469f881514323ff2906af410906319e348c0f0726d8b", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "39b6f2a0e6128f05750915b9766bf1ce51bdbd387eb2a22fe123f1b4256a4852", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "a92f2a79a40fc1659e33a221993a09ad906568251344b9add397fd9c9298dfb7", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "d68fd726507da0d3d9a09e43c671d9332cb148204da0ed01d7c389dfe7d3922d", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "4b6611e58bb759de24f846cf0d078d7984f844f9eb982ffd70d7e71d14051c7e", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "9aa559fcaf6d6adfe6e5446278b511bc32451ff191babbb521ad4d3a285a6ea3", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "39c1582911c80be5ccbe6929a6cefc3e4341299b5761f0afd38d0727c9325830", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "b7d56df446c2da994b72d3c811cd04ff8a8546097aeccce88e4f6ddecff5867b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "cfca63fc1fc4eb8f69afe06f1473daff603cb2953ff3f85d2791f8e3ce2c2a50", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "730984800b1448ce50780f1a5b6544aa949c4025ab77611f857048134c26d572", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "39d92c2e18c49b312625e7c4bb294dfe409be7b459380427efabe8bc50d2bece", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "09668127dbd80a9fc3c582c8ff165691069cf6a365b47dba18d78505470986d8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "cc6f3a34bdda4f6d794f629f5b66c602946b1164345390eb1f0dd21285be42ab", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "84d55960ecf03d1bd3221e9095c1ae63bc1f05b808d784b836a434b0b543e395", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "beb7ebb94c43f774a93c0cdcd7e8bb42da5312672730c362f4d11f89c69a404d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "fcdd16973f2139c9e3f7384452ad9ab10d4265371370753a1ca5e16bce53f241", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "1c99c976f741a99dec95655a462f4129ca303e4cb3c7ee3302107071aec7fb5f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e8bef8096d980201c5bee62ec9a56d8d0689488f416252154f118135a6d1f743", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "dc0a78777d7cbf5466b24fa203ac59ecd13b8fedfce37d81c46d9d077098f80a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "ad21faa8c7e426e60e16f77d054de7569430d3c29e61895e7f13ad20de57b232", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "26a331d6591dfc8985042a776ce256c509bf4cf6e423736bc050cc1f2d6525eb", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "0487655b74b67105e7ea4f4ab5d207fcf6c1b78fd78f4d638ec98499909ac94c", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "b304c159b9355cc0c5bb538c3f91e660f79691d66685cd75d39bb6e5b1bab030", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "50057e268157e460ff983af6b34efb53366d18866910690342df08831927708b", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "bc1d02414675727fe84690b3d4872d9cb306b8fe3effde78a1ae46da7b0b9107", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6f6d095a8e36ffc8a6503c9f032c7b53b0af85bc46e8883edbfa5c8662046dad", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "3debe9e90a03620f4a0a32a9f5bea4e9f0faf55a74013248fbe4fac493dda52c", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "8956f0a39b6ca444d74c2db4d62a90893ab645ebf424bc031ee10084aa881615", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "04257a336df0834272c9672c67fa814c676e473ff479cec89d994911335228f0", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "4786cf517276b67158fe60d677d716edfc869ae46acf3f82133c90054dfe9f3f", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "36c4506f1c7a343c242b02e929640abc58ff6e74014f09d9c4465adfd79639ac", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7258ab75a4ba902aba73710da465a500e207c358acb01cd9638bcdab682b7498", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "e34b03fd3d0612820b907bda92397b765d475e78ee2c317c6c1195471a49b500", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "60d64a3f2c2ba0ea6615a31ea1256acde0d4fb8f2e38432bd9c923e6892051a4", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "69722d0b8923d3f54a8a44675b07cea464bb6cd4a9c5f60498a4c0e9188ad6d6", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "a5ba042de87ec4dfcb53d1c5783cbabbff1fd82d8b60007ff7bc6c79a0c38bea", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ede35172557d9a8855c78c480b377242a527acbdbdf8e381a40a1544634739f3", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "1cc2d7d0bfabb1a92266eb9332cd995399051b5bafa57d99a5c091dbe78268cc", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "2c285b4e36d346d8c96abd22dcf9c8c7012bf04d0f36072c3047bf3bfa3eb57b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "ffa2fdcc3d35bb46aa5de593bfe4a8ce0fbe3ea2a470d11600f3a72da164acb6", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "9ccc6ce50ec1c6b0d4182089c955b5b2ffdd979ed8b2a6fa3c443f699e4c8989", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "536b2cba7760c1d61ef205f9e9f463246f3ab3c8e04135204422e73590751e84", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "87684618673fecb895c5059c446c3447adbcdb54489c61edbb12162875773785", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "cf6df82f22cf3c4ae5f21e943da2614e9fb813ef5fb20e56161499700544d61e", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "dcba82c9e47860aeeebfd66e02d0b46589196116e1330cf1bc3142480ddd8b9b", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d3c671314d15365df6e821550f53e33d67bccb59e10ff628596130dd3d64eb0b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "c0c1e4118dc07f17da3db91382b9e01db30816c3aed2d8a73a7b51179fba5d90", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "5facb7ce0fc2f941631b67b5590ed93a38bb3685c9213f477f7f6fe56d8e8959", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "c32f1cee463cc7ac33b325024e32d36f3f7723e057ae53227cf4cab92739f5ba", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "a475bc175bdc86f425fa5a6ddab594d7f1130df3bc3147b281eb7439e693feff", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "c6d5e5d22e655501257f71d6ac9d55ff61b10566f077866bf569fc0272452170", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "aecc054ef3526e0f568e4765658c92a10a4deb015949310c6a657374856ab2d1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "395e4ac1745b0b497fee38156df2ea0d7ba85b58a0f25e2a9d1733d6f20ec049", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "7010987c2d873b99e28942272e088431ea630008671c5c7c39b82bd0bfcd7983", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "3b4ada70f9282afcb194421d74d8e5caf796c5dbedbabd3be50bbea304e39856", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7f2455d56c80f6859518ec6bdd8eb7eec37d203cdacd38391c335278f88e70a9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "837970e9247a6439b7883c3da113372e7ace1034ba640c160fc62dda3b33f6c9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "9aa791f3dd16bb302c1e9926aac862170117d99502266d949b5ccb120fe1cb55", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c43b205c9fc5994228c52363a207be5f1d800453154f647f867e6f8bf6c5100b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "2c01283fa3676238918245c74e9cade24182dd541cff77034a222cc0849a74c8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "3c1f53a44d076a5f7396c802e4b790b32ea708d155a3c87f4493aac302eff542", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "60de8c0ec8256877357cdc107aa2aa553b254b7e9747aa55e1f8b340cb2a7438", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "aefe8150dbdea0c86b53963440c305bdb475e4d28858cbdc4ec18eb35242fdf7", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "be3bc1296a83ec7aca60532f3bb29d80d3f8905d7b0ef3b36915c0e85c545d0d", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "40a079b2608c65e91e0202182b8048e7b36eb7820db78bf2446362d2b995c549", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b772ada05f10f0a739644ecacd9cee246b643f2682c914ea7f61f28dc5dcdb34", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "8d629fafed933b2edfbeb0053c961a2c40f25a2983ea50f6a9fdfa97b49a60c8", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "4a34b90d092b4ac8349853624e7a754ce6115c92e8c7ebc346cc395bceea28c5", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "301346ccc9f1ea26ad4b42b099a430fc51dedab064a41d465966920b1b741041", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "cfc09f776ab90382bae739ce80fc28732886a22f1eedf3204795bb9e47ffe0cd", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "9881f810bfcdc9146cd446399172520a322f9eb00ecb9d6b0727855a538c1ca7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7dbbf67256f6083a261a235e2616f2d88ae556a2c540331e9a4c060dec2b0f29", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "570566612b11a3f4bffecd3a05813c874770a5a1baed005f250b087e3a7f7452", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b48c2f184319a1b93bbf1d3d7f9a6297bb719795958bacf41595d2624b235183", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "9172997df81567dc220125f054d9542d74ed1f7db8ef16b47e4a5590ac092f00", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "5c3e4a3a49aee7dc8162ce2f3952274b403415ff13d78da980de4114da2188b4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "8f1b9285ce91070daf4c6d6ca3f28cebf2a5af26fd4b315ae4eb3dfbc4d340a7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e83b3211abea0e1f4db750517af7ae9891041f62a2093042eda759037be8005e", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "6bbd1bacfd429d237ecd60d4e6a71d3e5ca798e919e2a8731538d08e27205726", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "656beddbdec085a5ae186cddf932277de1c498b79663372566cd2d8fa6f76279", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "2425bfa9ac4a524e5db896d890d8c207b699c414d40a92b5d227c96ed4d64f8c", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "42daca5358a07d2d99168bc5c8139ff5fcd744725b919d85ee42536ee78b0af5", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "32b0fa95393d9b9e2c37bde9d446f06fb6811cf20684d45d5b8d2403e550e338", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "4ecd9d10d6845799df3666865028f05715ccbd953275b32b77882b973caa79c6", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ebae0aa5831348a59aea382e24125023ec4b608e945ced3fd43ed91ea053ea17", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "cf0115a8556e65b20c40476f7ce201fcc5eb0f50b2cbd75894afd78ed6ebb988", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "99744cc8a662f8277f43c5f902a3102ea47a71c89ac5edf80b41d96665d62753", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d75c3727803b4a7ffeadbe1044ef6b8ea0d58f97cad86db25246daa68c69e78b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "a9cd524aa5d4b3379accd553bf5873b4052a264b9931090767a44bb15bd7bc81", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "44521617668ba67f35214d1db06299692ba751ae67d73ad2574af459e607ed6d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "166f95dbda1c5a5d03007fabed43cbe01c712807bd8aae52a19a420ce25e9b0e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "ae2350d5768f413e0d75a81246ea823d03619bd6792e700883073d2ac48d3b57", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "33477000289efb3a02173a8d087a52a43eb66030e92347987121f7e48ea3fd77", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "02d8093bcaa633c03a49dc8f12102c3212c1dfe6c768efd2cad51722afb10e0b", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "6180885b09ca507c7e50a420ce704931169f91347fe2c1c78f31c428179180ce", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "cd7d30dea8dac08d3521f8e8fdc97283504a8a4f71b2f720202b3a21b263bb7d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "b479291a7ed7d32cea45ebc74d3bb83a5a5814db5ba71d5cae5d249b703e14fe", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "afb6902a5a2d79212ebd3b8f59069f8610875862ddc4d18395d5cd2edebd2f02", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "333123b2431e50c5f2889bb5542fa1e76f62423abadee9317421591336542c47", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "27975bf4e5a88238ee4413b1b259ddd31fc03a7730c2f6c0589148a34121eee7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "c9c3be7b929e3ca4d3a695240dec463a86ce65d1c9bcae69ac07278d31d32bdc", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "1374fe25e4547181231fa337fefbcbdbea2f29b76d7fbffa0cb41c1d4871bdd4", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "25fc6ce1112cea8052745152b4a3886d426dbfdb39e457e47e77dbb9fd71d5d2", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "657543a4ab2cd6fdc26b4cdd3943aec5160b4a10dc3753aa42bbc2f0fc3e2e28", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "2ac1cb159af8564d737dceb2e3feb6d60e94b458ad67f5606ab7853fcdcafaeb", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "80a9fb26469db61765d08c332cdc34f671c355a01054c5c01b05a155ccb9f0fc", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "388111579abe178ecd5b7a3d528f17899ff260fbbba59442db303bfb8e812b12", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "cc35c0fde9756eb38d9cbc3711cbb1dd2f3f305feb92bd2b107a9965587d1fa9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "ae9b2fad17ea90b1f0ba5a6b511bf909b025ff331d2198c22bab60ea483e5b27", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "18ef858ced4d85cc857c0c7d2e196219229c5992c33b34575f826d3cd1e5fe47", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "f4aab63df659d9d2921ab4ce51edb93453bdc10b535448d18fe68f87b233dae7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "09aff9b4682ee052b01ffd1951a1a7eec41cb9f8be2f96e6b024d48513a6a5d8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "59b2ffb520f9b29e37089ac663a3a2a2d8ff69acbf9b372ad5f409da6f8bae51", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "4017e1ad9480c12b4d138041f777aa77e3b3925414882a65296ba22252a1f826", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "3c5d15b744f1d4c3260836e09bd23889b744eb2af5ca69f20f646220bef1b1c2", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "03aad604b8ee998e059c9acac26ea20a00c482286a4cf6466aff0dbf9c8f338b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "2cbdfb5621d756e321225cfaba61b8507fd72ab8a9435386f91526a59c5ac123", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "173334f402c053c13d350be1f5f956c04690ac4f8a5ccaacfccde1bc741b88cb", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c4f4a97d3b81f17158a69899232ac91809278124353ef84c3d6f26967416ce46", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "a29d0b2f14ba6aa108ee90d10346a08d9834d564b6a24bc54a44656e3bf308c3", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "7df364efa1389efd2edea2a8db56fafb0b2e9cf0f2542d92a04376bfbebd116f", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "a32a83e4c52d36ab827aca88e2a2c88a0817013794b609c3caf5dba23b9be974", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "20a217ca6ee34ceae0ea802a2ae26984107f994cfe697360f13d242d23ff4a79", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "5c91017851ae38cd6228e3d7385fb312c978a8db177d56374a8b9d79c46341d9", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "b2fc09655d7b25ef32f19d25d8a1efaa31e253f0d2c599b67f36408bd657899e", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "c50660eeac30efe9a97d34518fa97adfd22855ed72c3e48140914a98a10e03d1", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "7df97fb592149ac73970683d55398e47f5526340cfb9c79ea134ec272df61ef5", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d2a6a579f2e9a4c58130a3a3d713e6c9911172a3b8e4423c8f4b21f5189791a4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "7915f30592f2de1bc13e756c6e471dd9663e80f2b4b6abb820e63c78f3c2645c", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e673d0521299ae942b0122a4ef8b061430d086d35a46ad99d37dbd7fdcea98ca", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "1c3ff2eaacaab91c6bf566961d8b8d0aefa106b23d6afde146f17ab1b47a39b7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d6c007e5f3e55f6456e4e5f70967ecc33ff5aa047cdf74c597d5718a5b7ccf2e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "9cedac34c053435844405cba2a0da30b18ae4631f2fc377a94e18b42609853b1", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "99afd82f0ef30938b64f80954091c96b18fd53cbda2aba0a66cb77ce9c75a0f6", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "ab34aae7b3e827715aad82a5c6791a0f80822b01b2c45ade85fac8ed357547d3", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "ba60248def919a10338cc06eb7b7ec48bd7c455846385d0cf2e6fd203e572284", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "2fc71b39291f064aff4059f4aa4c58010ba9e939fda912b93a4dc2fec8a0e70d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "ec99c13c18c2700f3750de1ce5e56125b24bcc1f0c2a256e15be82bd1bd3835c", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "2fa6ac41de6d1a3850198cb66363e30bdf28dfeb57c315bc581d07caeff26803", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "f17586046037798e72e9a0de63b54089984fca8f1939a10f83c29635c10bb88f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "208252dcf8792cbb6550530c1af68eabf45576dbcbfb27ae22091ff246fce9ae", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "51d282001d636d7700c0d72d14f6dc897408761f5dd67e4051f70db362bbf698", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "f6bea628a4f248d73e6e4f229fd215a57f00a4ed025e545baf22550378fa96c9", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "230dbccb762b9b0b30fba26cbcb18a78139fa18eba237edb4593881302d36f1f", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "77d6650424109cd393e938b7d48892273d699fb3e333a7d6141478c9e6fd7128", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "307c60a53748afd1ced85698a53d989cb3bbb75c49266834bdd737163813ed05", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e7cd58d63bc6d8101726da5fbe233103a7b29c08378a8fd7494f39055af0b98e", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "93782fc1b484149bb16c9a844fbe9e6c207e8ce8f241c4b0f8b4ec76c036c02b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "d8a1813b1dabc75118e4d21dc881c02fe301e80f3b6d84010767286590b7f46a", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e13917ec3173ff1019880f31683002d45210f029ce3dc7a576d950be2703f25e", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "0ad6bb847c9219826a6573fd6a2a2caee468c5eac9af8928de49dc6c7605f862", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e6659b1da1243d8992ee2f4bd768caaf532a2d9b24716749bc618f65833500d0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "8a8d062a9dee9ab0998601efa32ca7061701d6bb1953d477f467c75e1446cded", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "2e18713cf2b55a4f5296ca32a5de73b5fecf00185b446caf6b504b4f77ce882c", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "128d98c84f3e935849e23c8d38a9927e72898054c3f98c82cd8a7ced6a7f18de", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "69be0d09ff0db8c9e4d95ec95289739a31914e3505efff6bba1624e378e72ec9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "dc2f3271514d1b60b254d0f528be5ec6adfac4ee377d050f6eaf9042631d59c8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "9e27711e453f1bc5a5ad1bdefd3dd117f8d653b5b43b22b4baa4b4fd9523c47e", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "7d490acb57542a729276cd1ce135273510489c4cbf61012635f945cafee75374", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "ee8b498c8c8c0763ca6eb4c7fb7630f3efe75ec39e874f1ad96f9d34b3dd1e4f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "425d7d76655d13d936bdcf8e794abca7400ebc8620150920a2ecd7ad49fda551", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "61a32e587cf9bb6f624ba8aa92941c70652b2965a2db1b410805e13ac9bc67c6", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "a8ccca2b399fbe33a1c705353fe96ca3b923050a8af659aadb06c95133b640e5", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "3aad77d8f338c4fe0b730e5a165f1cbd0cc4b4a61c011bf7affcae0bce41747c", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "812f48b9537a222222701b9e41d9f027db22477537e8cf141fc11fa34db4d059", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "83e88a2ad22c6aebe70d1ab25f897d554f4853001ddab5b6b2a0400a36b064f0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "7b99645177b7c61e4cf54c20db0ea402557f4545650261cb2304c16b74eb33f6", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "5c382ce1c61c61eaa1a14818bca2aeb6c93e1216601f38d14673c2c1acc9bcb7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "d1b923be40891f9e81899724bf019065521b41165331adffe5e14b66280cfbe5", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c0a279019311b4c5f3429f6a521598ac6bc5cea03670452cd29736a1b38e19e2", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "22fcb9c7c20253c8938f0ee69440879d55e5fa0424f7c4de6ef33b89e42237d5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "ff828c14112b74a6c6f5ffcc9274e6127a3e5c4b1ac26caef377f501681a6373", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "f9de993f51670418f40ad12948ff080286f0c0e1292f8189911b01ddbf5255c0", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d104f4b6077e494b1c5e4dd9456de9fc5ed9aeeb7bc17f61f713b85f35b38e6d", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "f5f8fcacd1440b994e4fbc473495f8e54ca20b620dba2b8b7b12848316bb657c", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0dadb924d802544628c7d2210cbcd43b91fd250f7dab2ffac1640fde31227b16", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "2afb5e0f07192923b6971cd1f6c76ba5d86bb3a2daa716ff3acf63b8e281a94d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "4bfcd7e8949182b68f2852ebaabb7a6ee9dfed86128ac27604958fce1bd63932", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "a4482cdf6110a4136ba9bd255068632e75089b18364415d39bfa5664719b6037", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "3d3d91d0623b5a8a2accd56dca260c7820cac8797e03706691961033f1ae1771", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "1d0127f6e0f9db07e827e5a659b277f68dc1fde488f5f9c913598603be9e9c4d", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "8799ba83fa649457d61b1698e839a0f0d1f21cd9963df0ba8a125261800831ea", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "734e405d1f63c5c617ba1db7f0917271ff54a5c02d3d451940d6a3b846a44298", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "7530d425e54d7bdc49221afbb65657af8a97c226b566fa55277772462d593a75", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "cb910336a4c0d90aa74b4d2794b264660bfe2c743de699dd24933a448acaac77", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "34a878d09852aa8abde06461af8aa6b556e4b2c81f0d21b15067862968edb840", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "c16b8ccb37d2d09212470a3fcd037fd6d0c97839f02f83368fcab73a568a3268", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "38f273906ecaece54cfa074a297b768548d93a920d33ffbed5e55b9facef5154", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "2beb06d7cfb8305582b5ab3d9cbf3b9aedd96e6019c8db6a9a60bef4b60703fe", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "d0bf9bac6de7a8203a31f46dbde9402e5021fb18aa8cacaac38a5ac78aba08e3", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2c538e391092347d054d33f05b2054b16763e1b400fbcd84ffcdd9163f6a6ef9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "506af01754a8013af669e9b1173b60c218dc437fca005a433ae0d2d260d588b8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "054ba44da560ca7de658efbba838f4c135da7bb48081743b3666008a9f306eca", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "6b89e42a433e8ffc9d4c70c9e8408fa3d0fa994beae178b3adae04e8f52a539e", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "1066e5da0be7574b245d9a7f8238dd38dcf475379bf478c4cf57882da605a345", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "c24ae4360dad963bdbfc74254dd86fe311deab60081e0cd4d8f4f31a905cf2ac", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d27cebc27112d251e3e4f6fb83768fded87906e1c9f8f9a12612829c59a8157b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "dddc0b5ebd88e07831404869d4e884e854527dc17750460ca298ace65dda812f", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "be7a3536be89142d39851f0e4033a822877e16c4fe104240bb6a1ed523642018", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "92084479a19946c87f9328d039480f2f98b6f46bee34c66e5c9c77ac2c73fe79", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "e9fcc958344101f43bc3d35dce9b2355f4dd96966c6af087227161ade792c714", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "dbd4a17f50139ba9529afc48fe05b709e8c319edaf9e77a9d85dd13658864748", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "3221d9a9ec76ab40a317c9fb439ad6893b389bd66ce0eac2ab55ad8422f2adb0", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "3fe18c53dfa1d990cfbf4c78ec8e5afda1b996e22cc3d849fa8c6183c12d6a5c", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "810c43b3fab908ceda301d2180e95182b6c07ec38f3ade3ad2605b6429d45e10", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6bfbafdb6cff475d46f9e54f0920913d645fdbf4c4d259147006ac2007e5faa9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "79773e76374ac1fc649172bf85808aaa26b054469b2598cbf1bef3cc7f83ccc7", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d2913ec8046f79839418d9cd7402dbd4342c88c72839a75a8f88052e99537100", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "8c8417a025c91e9d244d704df198c07e6e289c0868dc37e68b7865c0af647573", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "4c77aa92ea425c2e2b5b1a6335a7c54f40729a03a02d624b101be113ab072d56", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "9f04caab67cabaaa19500c97a7d21b7727e3b418417748eb86283155dd1c814d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "472e9845031c04aee34068abcb5ed587c3162f333f97cf981797fc30f263f39d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ab5531d133a4d0a1eb7013859d81a500c63b93368016956607545cfc7bf76f8f", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "7d0ccd307dfaacd10459055b62aef2316fab1cc7ce3eb2bb78b0f625490b659b", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "47b6f2a6e012c8514e3c16d7cc1f7043d39634a3e2b21fc1e2e5a930cb87eada", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b8d10a8e55d6ed2ff39531bca7ace35c8846dc38a55996754c0f8ee85dd9ab25", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "9240bc6f89ba88012c70bd38eb4ea91b5d7ed994d4718c5e8332f4c82622b6f0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "bf55a0f4f7796216618f40eb521cd9547e91be2d1ed9eb743f80f2e01345d70e", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "7251e79b3bb77cbc917456e9edc8ba3c509607ac78922aecb01c6098ad8b1ef3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "4c55fc5828e15e958dce5e5b4b9022f81b0ba3e2a6c8b9d0e918409e13fadfb1", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "e9de104584a344c8dd011bd2c3c267ab185a9c40bc1e46d6036d4d791367b948", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "288e3c30c44f576b6b86a821d75800fbb184ca5db86de77e3b78a8057c5f18cb" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "0ac13e8dfa812f09151d2efde6ad3904466e1c46682f550a81ec629a1d91300f", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 29000, + "training_contract_sha256": "c32f279d672acfa2d0ce7c6cad11410b145c081a7716c1365909274e0c11b0a2" + } +} diff --git a/models/ppo1-diayn500/0/_CHECKPOINT_METADATA b/models/ppo1-diayn500/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..d2a01c0f98d4ecac7ffaaf2c6fec3b3fc1975e42 --- /dev/null +++ b/models/ppo1-diayn500/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786263789116544870, "commit_timestamp_nsecs": 1786263789829495593, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-diayn500/0/default/_METADATA b/models/ppo1-diayn500/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-diayn500/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-diayn500/0/default/_sharding b/models/ppo1-diayn500/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-diayn500/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-diayn500/0/default/array_metadatas/process_0 b/models/ppo1-diayn500/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-diayn500/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-diayn500/0/default/d/a1433643c777305b17c52c5376bc0c57 b/models/ppo1-diayn500/0/default/d/a1433643c777305b17c52c5376bc0c57 new file mode 100644 index 0000000000000000000000000000000000000000..19839ba2c1835b6a1ed223e082cd52a26b90018f --- /dev/null +++ b/models/ppo1-diayn500/0/default/d/a1433643c777305b17c52c5376bc0c57 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:553b738b9a783d47f1f40444b7dda8b079ffcacc7547d44f8b642eea1f6541c6 +size 112878 diff --git a/models/ppo1-diayn500/0/default/manifest.ocdbt b/models/ppo1-diayn500/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..b3899499056d38b2ecb486bc76d0e3fc3ad2fd21 Binary files /dev/null and b/models/ppo1-diayn500/0/default/manifest.ocdbt differ diff --git a/models/ppo1-diayn500/0/default/ocdbt.process_0/d/265ace8b1373aa11596e3e75feedf482 b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/265ace8b1373aa11596e3e75feedf482 new file mode 100644 index 0000000000000000000000000000000000000000..28f66b5f6b7d38a9023571813f1b6afbbecb835a --- /dev/null +++ b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/265ace8b1373aa11596e3e75feedf482 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:da52e57c861351fdca66cfdadcee2e2a8cad29af5bfd5dcec887a497fb5a489a +size 72097792 diff --git a/models/ppo1-diayn500/0/default/ocdbt.process_0/d/3fc5aabae3972fda1850bdd8d2de913c b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/3fc5aabae3972fda1850bdd8d2de913c new file mode 100644 index 0000000000000000000000000000000000000000..e2a24cdaf59697cc0b1807bf9d5a88094989cd09 Binary files /dev/null and b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/3fc5aabae3972fda1850bdd8d2de913c differ diff --git a/models/ppo1-diayn500/0/default/ocdbt.process_0/d/860e7ec4b2f6d773e351ef0565846806 b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/860e7ec4b2f6d773e351ef0565846806 new file mode 100644 index 0000000000000000000000000000000000000000..093b5934caf31def1471feb25d4984ca97685983 Binary files /dev/null and b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/860e7ec4b2f6d773e351ef0565846806 differ diff --git a/models/ppo1-diayn500/0/default/ocdbt.process_0/d/8ac165a32779d4aed6fc4892996b8dc3 b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/8ac165a32779d4aed6fc4892996b8dc3 new file mode 100644 index 0000000000000000000000000000000000000000..50ea24a16600413ee19ea7c00cdfdbadaeb30ef8 --- /dev/null +++ b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/8ac165a32779d4aed6fc4892996b8dc3 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9e02525b32a92c81a51d5886ecf9ddec178eec06d704876cba398b31864ec50c +size 17539072 diff --git a/models/ppo1-diayn500/0/default/ocdbt.process_0/d/bd998024e8a75a5d0f62ffabf3051b3c b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/bd998024e8a75a5d0f62ffabf3051b3c new file mode 100644 index 0000000000000000000000000000000000000000..a3f6243469d2a028d6af2536ad32793e2462137a --- /dev/null +++ b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/bd998024e8a75a5d0f62ffabf3051b3c @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:36e0443daad45d7d4d78ec806670ae5d38a5a5c2efe29056c47869e0ee38f21e +size 489701376 diff --git a/models/ppo1-diayn500/0/default/ocdbt.process_0/d/ee12ad97855496c5a52dce9b8e4e4e4f b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/ee12ad97855496c5a52dce9b8e4e4e4f new file mode 100644 index 0000000000000000000000000000000000000000..a7b9c8ba17eb23e154fcb6800afe23b886e6d183 Binary files /dev/null and b/models/ppo1-diayn500/0/default/ocdbt.process_0/d/ee12ad97855496c5a52dce9b8e4e4e4f differ diff --git a/models/ppo1-diayn500/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-diayn500/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..89a832a5c05776e9e5185efbe5989c9f0f0cd6e0 Binary files /dev/null and b/models/ppo1-diayn500/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-diayn500/checkpoint_manifest.json b/models/ppo1-diayn500/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..0eb6feb04fb9d9fdbc7899b49232edfd63842b1f --- /dev/null +++ b/models/ppo1-diayn500/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 34000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "a32d5db947a95b2d6ed75c2a54481b34cdbae850d0831552d8252cabb99c67a9", + "runtime_config_sha256": "a32d5db947a95b2d6ed75c2a54481b34cdbae850d0831552d8252cabb99c67a9", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "a32d5db947a95b2d6ed75c2a54481b34cdbae850d0831552d8252cabb99c67a9", + "schema_version": 1, + "training_contract_sha256": "b56a43f34f8cb966acac8f0fff7f69019f28272673ba3583c72bcdf87f802860" + }, + "sha256": "bd3bf439b1b316ad42548d42a5ddd605e1dfec6f6a6cf6ea16358223ecae41f3" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "25c43557d9dbc28bc2559496fd25904836b81f94a70f2f9c3701c7c7daabe665", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "fec97af9f57adfd3a2c28ebceff4b4b1f48e2170b45c60271efc84d846edd7c0", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "fe10230aec9ac7415e66a5c37951dcadfeff033d5cf33eae63eeaa9fd2bb8b00", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "3f7e5bf66198c278b78e1866e1641366b09e377ee19e05e6fa9c749699ed6e6e", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "ddc0125c5a44ff655967df43bd4ac8fed39d249bd2c14aa77eedc45ba9083b2c", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "6daceb557fff23e5ba9643e5e1f19c8c4e754b35b471b60634dab7b01723536e", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "d97e6e09de6165eb89240ecbfc0b1f767d2dbe80f6abb4f2eaf7c1c6daa9271f", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "eeff356a38f1317b1872d72f4563628943b8bb8816944deac80bf19d8fa67bd3", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "4e0eb88b4619482484d682cb603d8c0165d41df99a24b5c84f075d32ca841426", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "880297b0bc71218db8e0921baf81ac017ae78c474e14092294dc7972f87215e4", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "2f5a397d4519669dc6ef0d00b136482edeaf769ddae25bccda543fc1937a0413", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "0732082f4c039a51c492751b56c6dc2d6785cff28818ae5842edf5e4d4fef215", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "1e3255e10477aabc205d072ed42bd91ab088401da0bb0eaa1a7cb104e4743377", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "3ffeb55d6249c0060d3129d278c02a666a8215c326c630e3707349a9c3d05b92", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "0c749e887a268d45a9275b5db2520dede330d79c5726fe3c1dfc986f80ed01d0", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "fb7c1b856c555cfea602e291224575e74a70769679dfc57ac8b35e5c71228959", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "52658ca4bc6bada5bc7de9c084189cd56ffdc35dcaf3979471de1db9f386df8b", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "1c10ae67f9b621fdd8f225eb97d173165553f0882f255efba3c816e3c2387151", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "25aad4c293b24c9826d0317ef1ce4128d1750de0c8e61f688e7bfeabbce3b794", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "06e6f6e546bfa2ad9532245051c057ed9e602a5ba9a638441316e9b4976387c6", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "02c449572eb36d7274bf4b074870204e3491e7ec178f56e66ba50979e33967ad", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "47328ae8b58450a456ad2844041adfc763d18307841c836c8034034e05ce0c37", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "8c58272b20b0c402f1035be8780f33c6e6a514c2d6164a27ac041f6fc2e369e4", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "a946f92ca39340623765cf6685016b04b904e9fd931905ceb40c2bfd9689330e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "b0510c22b031245697ef7f0a11e52d7d82b00b1f83d33b3d51ca4bb2829c469e", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "d782255a3fb33da483fc9a8aa94af6133df76299bee52670e1dfdae36238cb02", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "ef8f80b7e1c2c935debec9268498df809bc518b36aa5c52fb88c4b126e560c09", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "9eb5fc902da8447776702e9a2e9472d5b986e865eaead2cc957e294aa9f9152f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "50caabb8b579b7928e4b4a1473869f408c55b5b99f11f7d264d4dff3880a54d3", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ac228ff84bb177d7fdb181871c9278eb0ebb25e8910fb55b1c3659330d3ee36c", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "4497203cf1ea74e31343f60aaba4e4134cd74fada4729ddf33c8a457cfc5ba14", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "90884738b06a273a1c2db6d23ea60a07362bd22c093163e4e3f0117edb715f85", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d3dc2cbce2981c3a877f1f297b06f5b76ee3407b4592dbfdd2b28ee171cbdf3f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "ae8f333823d234f7f986344ddcb5580bdec6562a456ef3c2f6a9b89e5ceadcc0", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "825a11dc825aeaecb1075fe396ac2e760a8bff99a08c4c7cda7244c538b16b52", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "ad3f44d0564b44348770e1da70c71109cfca51b801bed6d10892bf18e2209b45", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "476c1fb4868c4831bbe0904040e0369778fe347944c076d436bb4d1c6d36d264", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "29f682ad32159be5d995b2ca3bc962021c2e7a38eaf6ebe566c094f56044d141", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "c63930a5ff2f7a62d723dcdeab2f19faa4cd9f9d638dd468339bd8c839d2054f", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "5feba0915df97a522a767914c0a764dbc3304ab5a51294e75f42129ea6911e4e", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d00c2a39b882b30411f950b4d9ab948ab9226fcf33dff54ac49251b9ca24c61a", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "dfa72c2fe42eb4263bfa4c89feb4fbeed695d8cb1c3e49b896d76548b1ec4bdd", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "4e5467f2a195d8b7847f70b2cfe3837e979456a069afc768d4bafbcd051a6ed4", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "76fbc8e03ae556809fa43a2a9e0a30455ce13a202586e3aca53b8a07c5a5b795", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "f928250b84cad61104a101f4765bc42e306446ed0364d1156c5d40549b829fb5", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "28ccce3ea11bb3ad10ce1756019c8fe82ed374c116257d297fe688b694a54265", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b8a0fff929fae8642669e6f3d324308be96cb94bceb44ff2cf8f290c7d98ab92", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b937e6e8bb379b3515fe90b73336ac8a79117f39bab9f3cca51774cb3f83edef", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "769d359446bcd4ab0291be7ac692aa48caf7352805bed7d186a3b39353f4209f", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "339b7dfad0c637c39c5f7bedf7658ed092ffab358907d6b0c65a03a8ebfe3f35", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "9a6eda386faa1ee8234de87f04dd2bfd6192c01c1d88f6d3dc4665b4087e3fde", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "5704256d731730643bb1098180d7bd64b0acaee8c9ae9792d2b90e0e7f5d8ea2", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "667b06cac31867fae3922d746142a7d4d218b96299b5569e0e7f5cd2155804d9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b2ad60a162e3a2b5c9507e8a30c205db63388f15f19c953c5785225977fad865", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "711998d13863f6b67a6ea498cacaec75d8998f3928a8319161e633d8c31f46c5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "b0e5e8e8e8a97c4c6f4d646327066d6148754f0546135d097651fcdcc761010c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "ebcc053ff186c239963e5bdd40e59a564da50a292ecbdbf7bd135d9d91c59cc9", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "3701744c1db6e9c5129b64b10b46e041e3c4e921c6d0ba58d7699e5950b7bf03", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "d69982f758708d0d49dba727cd9441e4261ff7e59540a86ef369fce47102f6b5", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "8d342d85fb3a8489259a7e2ebc88640ea26514d0f5d4b76602653e0cf86f123c", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "8b969ee25e7048a44cb2c2a0dc2bcf9119f4659288655d5571093a802eb81db4", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "a2b41a045b56d37d69e092310a6dd1c0aa9a839482091cb6af9dc14d0817fc1f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "425b69f28de917f0c7b9c668f9f72ee8210df36873a1ca76f5208774d252a391", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "421d233575cd5b90a3b38009a531aa80e87b3b45f04df0e9f0f44ce0399d08b0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "2b09454ae37640d975ce05882f01811ce273cbb61ebcd3803ad275cda43430fc", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "ab90e2d559093dac015c01d99d35056c7eb776ed86ce2b2af31ff1e3a5d181b7", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "26203e3a5c144a9db980b77076e1d4d878a77cae06bb6ba63b1d97ed86f48c56", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "6ecec1fc0af5a69181eb2baf592a94713e90a3f8c826f9a1440c7a00f0ea171c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "7870fb784be2d10445af597ef9fb7d894d0e7e0a8ed2304f0c7ba906e615270f", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "8d0bab146b7ebbddda780ce03bf764fde0636e6a148835b1150d9690fc065f19", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "8f951b692b3ba266435cb6cb09e135047ef4887a47dcab6c171f76c8bd96598d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "edadc977118b0e75065c2b565379fee0ef53df25e0c4b8632640821e6502cf99", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "bf059ee571cab9a61c20c5fb51148609da501a0c26f070cdb17944ba1e9572ec", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "75b5bbff5a35576995449b4e528ec61c0fea5b1960818cc38b3811283522f9f6", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "a518cfee4eb70c0a46718891db69c4df22e899662e614155ecd24d093ba106f4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "c809e19e3c67036972410b2af6429fe807d6c35113ab3b863a332dd385883ebd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "78f3bf10f364f68e05adeec5990173b5412e1e69eea723be66ab6f926ca05910", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "f7abd55dfca12e93c67469449c49d14f45b0041ba9218698846e4eb47da6d2dc", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "be4f929e31ecb3be0d7ac7d74d8753f57f9fa0b8043cc91b8ff23a471b109db0", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "27d0cb56984d7a82f0c0e63e791d34bbbf1c4ff300682755bfa032b600a5a463", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f7adc23321c7016d0ed8036824bd23155c3065070e1e8dab52d4c9b2e2c9448c", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "e67e9f04c7578b7aa2d64ff9be71ff8a682539325f8ea961f5a1e72969a35f32", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e49c37392f9c022391a0bdf562574cd511f9c391c8a2998c8c06ec2dc7199931", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "b4248a0549a4833dabec336c161a368072ac0a4526eec4b5b9d90a136aa9fb20", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "cc0c73f50ad1a89654c3f1ba5aa5ca67b9680bdb5f09f5284f16dcc40a3ccc8d", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "4d5d6649ab1aa0ce8bf98203bc9256b8f3f7c3feea64c60da1b29e86ae224547", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "aa1c6310ad03c2920d8c9d8117027c9be372479801a78dc9fc01b72dd9f0bbc1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "5489a1901b39d8e038988a860e844fb176c0831b502b5239fda3a34e9362af3c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c0949669a2f098c9f9366121be17f4be55c3bb87f917d3f2b4e656a5faf5027c", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "02e717d7910256698af83122709086e61944fb1e2669ffc25d71406434170fa3", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d1a747fb56d5941c76335921d595d3b9892b795c28738cb877ff544ee2e09575", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "a1fdadc8d56604d800617d77fbdae80f871cada6a5586db4a09333ae0bfc2c55", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "72225bc4f3c76c2b5ecb3247f35640428e5e7215adc8a9864172fc30315a8d58", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "d934404bf5d68c35507ee2e85ea496faefb384d2991c24e0bf122457940f7263", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d061acff3c8370c5441501b89d1d9dd5b55d243597a0785241516d795bad3f60", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9802afe623426a5c1cf83cf4b8c9e1577d30b7e56da9410668037cdf9f982b4d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "d0cfc546c8108555b2de654da8dd92237da99957a846456a718fa54d3c1aece8", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "66b8d37de7ae3221cdbc38e42e087a0bb6c587ad634fc8ad2424073b3ae2da78", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "7a16aff6a828b203b82721fbc10e9deeefc383fa15eca3cdd34b64ce3d4f72a6", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "b20e5311ccba15dffd71247cd22566e70222b90b1d1abd047c9553b5a552d266", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "6e0759c235901c40be69f21c93727eecdf0216d9b3e94dde50a6542d8f4b3e20", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "9982181084952eeceb564d1cd29c17c174133143ebdcf8a0b126aa29d2037de0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "628d868057f6e9650c3e55a54c58cc36ea969f995f2b158332208be79335195f", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "88e2e746ea3d227a95c7c39aa5e303f78d4ca2d7f15526f6fbc2aaf203b84457", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5c966da5dddbd3b09681c0f9f17971253052c6b011792a5ed43768c2ade3d083", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "37b150b4c52ad7a20d3654e0ec079a468d690e5c809fac2ed551c2cad3970906", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "412ae06ff27992c902666e06c4a341a3746910feb2a176b9e09056ea6f62af74", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "69293b88fa329c666014d42112f94ab659f26fd7c6f6a8ec090d3af429ebb1b1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "0a41d46bf01efce64a233ea0a27271b47e5f08f48f90daf845a67acd362c5f60", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "2f2ad0934170e08b101905f0dc991a28dd184b7cb97988c77df41b2e83e04bee", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "afe57ecca7d96d761fce0e10026099ba7c682f2cdbe41dfd0e03cff5e5e1b79d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "1e46a2a8917c89bcd5e538f100e3356e2c89be0f5adaaa88c7b31f877e281227", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "1d0bc04f7b4f8da052266fd4dd90f0c3e623a789c81bdabc484973ac74b3fcbf", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "36339fa4d754d8f9f19c570565314d42bd6027a9bd8589d4ea1d7f1819580eed", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "712181df5c6d9c0a2a5a325dab42aa3f3f53b0d2cff4c3ed663718dd41c96760", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9ba0ba47a595d71961173f79de831950f9747932834e90608048ee7aa7ccdfe0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "0ed439c21fe9539de14d333a03b71e4ac66b61725312a6df925ab45602d1dc07", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "bd472e5dbdb59527d94ccd8b60720005e9cd7c8764f16804b15e7b6bde033086", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "81849fed7d812419066579a2807017fc462bd9a4e966e75461194882b5feb4e4", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "c6f6e6988878fe551e3786cc7bb359ab5bac3385732440fdc49a1a8ec5428a34", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "96f3d24d51cbdc6ff9b4bafc9eb27d047f5419b8c5c4ec55da88d679d53efe12", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d4361a634e8cfb2e39b6b3d64bf3a2c3e41913964540e0f9e522011acf7c6dd4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f194081895d8a4d43fdfe44f6fa2896327ff5dd4b0508a6a941766aed97d94f7", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "185f1b5aff7f56bfb363383c9353788e4600ae6f8eef023f8dc408369d6ab960", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "7639281e22aa720929f619911327a2e125f2eb1e6f596a0fba90cef8b3469099", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e41377d84e252bcfb5b91619a391345061e4f32abab3289918ffb2fc1ba88dcb", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "623c907517d8f44edfc1e2060e9193b57e68fdde69943b967e38252b05a8d7a4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "90275eb075ede9a008b324242cec003b626dbe41bbe1f7b73fd4d2f86fea87b3", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "97fedae1b7f5ffb2ba3f493d48210269d7445c4f2c752ce9a7eb1380491e844d", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "373f9a960743f97e970989252c7bc672cf1494cf2e1eced75baed8d687d3100d", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "fd3434d1b49adfb8562d0180f8b606ec7283f9b37c89926f6ca2c6169a2f7ebb", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "f97ad644259c8350a7190dd3354d4ad539b814f0894501d6939600df2dccb9ff", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "99643654fe92eaffd2269cf3b166677532abf10ff55a6cab9fbffba2467062fa", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "2ca4c737b80e047f680d004bc3b2992fcd63900a108277ac280b590531286ba7", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "70bf92e98915d5c1cb8ff2a8627faab9b84f35c2509e414039b6906cad83706a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "76a073a23c8a90cfd7b382b7ff0545fcac341aa76dbd631056738188aa7447c1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "8ed2ed40857f152e60ddd9f4963e8fed4f04484294e8a33dcd09d1e67472e987", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "2dbf491c63e6398341f49f58f7c6e0aa83e65dccadb047b91d92bdf214bccf98", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "0a157cc84e01cbab0d341462f091a3e2133c54a9b36cdc66171c6f738d80aac4", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "1eb59214ecd5fca68a2487c6782ba51ee03fec7367802c7493da446c423c881b", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "762bc75f5dee829194d6954397aa47f508b78ae02c95e8f61f5a5fd7bf984443", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "45cb6046373fbb20f8620dda42a7862dfd9aa7d0e20d3c3726063b3871c50595", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "9520472460293f33f64516fd6c724204b5f5720a60b48df7bee18cdbf9ac0a67", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "82285ec395068f6e1b3bc25506ebbdbe953fd54fc2f38c54c484c3d5d36ad1b0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "2bf25d2455f4d98450d8a0c70df78a092ea5954550134c54828ef59bc6b1145c", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "b35c1cd9cb21af1e96bd6df55b0a42c60ee90655d7b24f3c6d6d691e21be440d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "e51796e3ad1677ef1c915b47d49eee194ebdb51179f95edc6ead9ef1599bcc1e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d67edb7678efbe193c731e93e0db3501c9c86457e0b35b7876e2f717801bfc28", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "10558a2b17d939f4ba35ea84332ad31846ce64a184dd4d4f84dc946d38bb94f4", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d5e0bf89a09731c20918baa435b59b048223d46ae0eb7205944cd7f8639ec6b7", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "dc4f43415a14381ffd137c0d6e919924465f58b929aa9a21f5bc6cd02e884b29", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "60af22c43a1d751ab20266ab91a1af7643c705975d306f5b1a6d94eb0238b1c4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "cbbb1121aa5b659a46e7e9135b92e7a39af23a6aeb7f8b00147f297f376317e0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "aa0b92f788cbbc99efb05c8125cfa4ed94899f9025774fd1e858dc28e3a162d2", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "0b20905bbae0f487be242d6261f3bc00d3b27f2a8d01e7a27ce67ae2ce25b4bd", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "8c54e4b4f6d9583d787c24268116bf945ac682a3165896e810fcd4b9ecf465fc", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "97b6ed4099b1c86788485f945b9c2c807db614c5caf4d1390fb6fa7836790b97", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "4e7d5dacb3ca80b3be0991021cbde12c1773eb953afedf1b1d95f829e5fd6133", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "211292d061b99165e69ac5ad55217e82f9f437a9d043a420aa65d559acaba19c", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "f80a6ffd6549e832ffdeb586f9b984dfba863faadf6368b2f38e32cfa7575f84", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "7072d8cb3e89fd132975da6c84f1ce4954141c0dc14435492e8a3ab52ae7acec", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "baa3a7299f2725477090bfff732b93f1fc1f2bb20e8ae4158c444730cab82e72", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "6ee0b86b49d501bbd12fadbd670339977c84e914177b2e0f29f67a80382c9b59", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "012e1db9016647542198b61dc6f3159f5bf39e9a79d03e1b43beb746c37d368a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "24427d3b2d9418d2930fb000bb26a206034db1b741be951299b63043f17574e6", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "bbb0e4ab3017a471aa6a3157492467eb2bb5995b4ae33ea33c89e51f96eb33fd", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a662ac9e73d3426cbb3abe6645ecb75c11af582b5bf41d001ca47be3c91f2156", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d578d6875438183710a9672dcf619d52371c1a4415e33a7f9306e5a4850959e3", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ee7c477415869f12350bf0565e81cbd279a3c9a64a673e0c1512b507e98d319d", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "f65ee415976d545d72626b83521ff900b2e8b15bc0ecfec96cfbec4a19735f7a", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "b2573d9c80e89721534fc980a8d690aa14f02cc2ca60ad2938a58924a352f5ee", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "58c510011f2ddd45fe5a7f853250d3b2ccc6d14d55871861f54a8ebdb6f7a3e5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "f0346b6d2d1d8c06fbf0cbd055ab06bb0cd481256a86521f5bb022c447e2eefb", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "7c63585491e49490f98e6ae56ce0e43c6a8156d960567622dd7dd37cbac7a4af", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "27c7302c45ff0a62a7d362da0c9b605906aec1772e92d0ccabb5494428d725a3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9bd48864be1878c79441769354cc61b7cd5ac527e47012b44d84f10d444dc47f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "5346cbe39cabd3d2e41a48b69438847eb5fbdb1b8b237df396edd2b0f57bb7d4", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "c3fb6333357268f8e13eeb1d00631f259dfb8547bd7c58f8e1a95a1fc3c1f347", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "bc61a287222e8ef61fd2ca7a7027135533fc46819be56e972caa9a0ebb2f77b6", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "bd00932114fcd706d942c378927922b7889c5db615fbce019d68fbba540e96d5", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "0acb34c40fd4d86c2fb3fe617aea5fa6781820a37eb73ce3d0fa131ed150090e", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "03e23f97383dd8987e575b0cadbab5741a2004075b408defe900754d7f6cd8cb", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "b00fb65fe7a546e14b00fd8bbed517ba1101d1105abde576a355114460540774", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "ed81fbd04bb1eb8d1077b767c9ab5729ba6b54d0e965e0feadcd1991084ba031", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "531910f31a70daa0d39fa76ebffd66061dd56d63d58540003ac1b448e1b79a0e", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "8ac15424ef83a91c17e976d3b33143d0c9bd7e1803f15b6be7e7634d701dad83", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "7916422b5cd22d29d2d83a502ca79562c76adf730502d5337f5f20ab0ecca87b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "4491b34d251144e0759920fe9a29f6b31b5c51aad9163b1f4ec4cdbf84135815", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "cb39a925acd54447e9acc088b4cd826d19c3ef510f7b6fd978bf0e202ca55de2", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "8d9e4dd96584ece3f81828606cde3df8e1ce090fe2a19527f4b0e3cc8e57e078", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "4143bf7c52118017e2a31659483f55a138d766cc8b03c86ee6445743e9406e87", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7ee462bde3743cb1e36e96a8685b892fcd284451123e70ea603d96d0c5e99266", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "935bad24cf9184ca48ca2c293f42d635fb919f14e568a091a043427271032d05", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b99a862b193da0b0c37c1dfc0cb41a2bc765642d7404efbd0ca762dd7f02a14e", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "a46f7efde3606b3cf671a866cff9598919acb2d2f3f10ac9c30253135a12ca26", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "0faf9f409e404e177bd254deeccdc21edd45dae51810f4c88a272e012df3885f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "fdda6e49b3ec810fd6920d9f147ef48cb815f4f4363d307be6a78e8a34410d39", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "cdc277f5f2c2be7611889f6ec04fdb31498fcc32e28f3a47c01ed14bc3772e8c", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "dc5dd7211fe4dde4499f6963a2ed76d52183e8ec7213055ec181017abef879d8", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "09dc4c2a853472f590b8cf216479e2ccb5a38becbfbf2a00a732e878609d097f", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "a1b28e845b3e366183f9744ff1d390eaaae7f8a6e14c512746f5b11344a0bcf3", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f08dff884a4ffc9b8542afa4d68213c86d877e3a2af60a66868910e517d4089f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "3a02248e0a0344e6206aec01888626104db490c78347de2140d66bd4e806bc13", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "b7c6b266726f09b96b4dde14394c66234e827d0690a8677472036c6961c8550d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "3cf6de5dbeedd0c0820ca85261ef042b6d3a62eb63e90fad1337ca2055857d7d", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "196e10417d41cc39e4e47c9f547218f408489d7965187758b22f4c850d47f7d6", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "ccd4ca65e8bbab7f27cae94f14c702be33680f78b517370860a2cbb63c951a53", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "4fe2a997b10929c9b9c0fe5271d1a9f6aec955be428eee04ca3135ae0be6cefd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "1f11e2da8f1710171d1ed78f113b979776955e931a3dfdef3994a61d301322bc", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "f4dbccb55a1902554d1002df66cedd157cb582aa86fd02fe9d3d92b6654de0ea", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "24a164ae410266d26acb35caa724429e7d49439897655ab640da04a2e763f159", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "61627b5cd8445032fda17d36870df83fa338ad02666b6ba3cd003f94c64aa01a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e85844df65d976109a195a289aa0ebe87c5638330e84a4af43fcb421a3d1f2ea", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "7f292656b7ad90790074af3b834837213dd47255f84b380c35f08d5bc28cd676", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "59673e1d2f79ba8f4e03c8cfead8d299a521b18988361b9d79821df5d422f18b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9b75b7f0227e6bc1ac798a9f95f841663c457154511ce5fc1e7e760eb5f0cada", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "3b14c16434d7fce9ab410c9003c9528f966d23addba73ceb09b8002de44b61f2", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "ce2a6cb882f5e8bf5ae30991fbe78c8d0b2f7915a3543802e8b4584edc4f119a", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "86c2b82672165d3df182a3a684461fb21064e37f73485a15958274bc3ad5cd8b", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "44235fa3c1b334b7eb02b9249f8e90bd7e0a88607fa89f40e3dad4f358a4962c", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "48d8a0c897fe5db425c2c270aeca21071308d19a4a635a6912e986ca85a53c96", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8923ba49930fd47a02d405feaed9d822a58c2f9e26df2d2d86892ea78f5e0f6b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e1b334d853f34f7e664864a2a245ea2d76ff349f851dee115b41d23370c9c80b", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "36d6d9fd0f30afec3113d8de5f99979006f7ffa49550910e03477f2246959379", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5672f997369a33f400c79458d72c5248d280603e52eff28fe8d7602073b9cf45", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "aacc113c87eeb1de6da4796ffcd5a54f9e2f77c0175ac5ad7c6c2c23ffe64db2", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b0712b0e4cf9d1f6c20e7596fbdf80c9d7e72d7514f948d77ef7b377765ad47c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "4440d1c13288a93f484fc0a3ab0bb0e2cfc8bac2145622b0c8730a33ccfca4b2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "775947d0acc5f82ccc7f6b02b436199da3a657b36931c08d45f28bf2647fc1f9", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "09679be68f954526074f82c90bf2b0e95ec451443bf056cb14a93358218e1714", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "3100fba882d264ed0a431587797aff83d52054ccd0cc6cced632dc71ed5e6874", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "78805a1dcdb2c918fa4bf0afdd634d225713c5dbf8bfee614a4d7107cdb66a75", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "26b5e3be37763c896702e645d8ce29cd759ad3a74e729e7126ec235b09e9f132", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "1348ef8e7718c4156c5afe7e337d1614300060703f4df4b56dcb5690b0ef77c3", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "cbda7df6165aca432d0ab096573634669b59fd3b9c350ecec0bd7afe2c618b27", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "a44120a4c8160d688cf006546046d9e7cb5441f7c7b1528edc6f79041e1a009f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "d2482704ce36738119fab877905e6bd0ea568a5250ebd5ce09a4ecf28bb07ffe", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "06ab6f577bf0ed2524add5a2abdc082b2f33ea9efd69e6aa6794d4591a061fc0", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "42bfbf797cbfa1778d19f1a431be7f8e73824812c5bad6225e03e7939a8a8d4b", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "78db7a9020ffa212ee982b6bd3afa6a46c42166bb3a609e59d98e47d714bff78", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "91e32897b8a362c5414dcf6aebf882dd471c9516e9ab0caacf85071cabee1af8", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "830f9d807143cd3fcaf8cdd111fe3242effdb5825e66e7ec17b1a052436dfe2f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "3678e71517ec027702d43dfb04389313803fc1caff6fb6821a606de647cb1da4", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "ac57fa95fb748888cdd66bd222a9082e2a6e4535433c526b3d521ae97883fcdd", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5c2277a9616553ebbb05ee06c82c4d6858d4b5f468906274aa32c56aea009215", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "de7f6728506e895771dc748a903c617748f011b212284fec29cd47d3dd44cd85", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b594d97cbe263d25280cb78a278aa12faf8658cdc8be918060526e3f3a67f827", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "6cbd13ebc654f3bf6ff5ac9de8637eb9d82fb5caf16eb5a599e0d908b3669fbd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "436014d719d8e711c3715665a63f79326f639640438007161f53256bda560467", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "6c09a5d8739b8acff52797b8e73b99e851685758d2cc447e1d2436e5f6cf6214", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "9feb93032ef40b4223f1ac419f57558e0814290ee0d4730d09f80ddf799fce14", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "82127a5b908d915b9d00852aa772a992f344e73c5b2d11c4a5070379ff15b425", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "f9bc971beacc4e246871a8d4a1389872e27ca78379f2c986b7ac76f05a571b57", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b95b01042de9e95f0755e9c2d2d6132311e45c1e178693e864336fb125289021", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "0233e9fedb2a2f4067be7b204e098dd97d2788763d8bad87140f8a1fdf6c775d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "0c33ada4c0c0459c97a278ce513855f85a194cb12fccfeea69ff05f570e924fb", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "151bb173b89c3ba740f63fd9aba6620b9855ba9f0228fdaadeb9863b62062292", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "ff46161bc43e2b4474350cb98a07d71639839fe36e93fff76537c7517e9134d1" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "7ee788b39b268dd06673600e94e53a07801343c6f3a864349fae31eea3badd33", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 34000, + "training_contract_sha256": "b56a43f34f8cb966acac8f0fff7f69019f28272673ba3583c72bcdf87f802860" + } +} diff --git a/models/ppo1-icm500/0/_CHECKPOINT_METADATA b/models/ppo1-icm500/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..d79dd49fef7e79c21c644ae01c5f54f259e70d5e --- /dev/null +++ b/models/ppo1-icm500/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786263789404404609, "commit_timestamp_nsecs": 1786263790081745749, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-icm500/0/default/_METADATA b/models/ppo1-icm500/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-icm500/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-icm500/0/default/_sharding b/models/ppo1-icm500/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-icm500/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-icm500/0/default/array_metadatas/process_0 b/models/ppo1-icm500/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-icm500/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-icm500/0/default/d/5dc7664f7ddd729e30fe48716300c09a b/models/ppo1-icm500/0/default/d/5dc7664f7ddd729e30fe48716300c09a new file mode 100644 index 0000000000000000000000000000000000000000..037a0b8828d78c29e3e03511b571687e79c5d2f4 --- /dev/null +++ b/models/ppo1-icm500/0/default/d/5dc7664f7ddd729e30fe48716300c09a @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:209f629934f74ee01d125d3783cabfe1e99ff327de24d6d8c034ed32b3b2cc40 +size 112942 diff --git a/models/ppo1-icm500/0/default/manifest.ocdbt b/models/ppo1-icm500/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..c2675e94581348afe635a37e71851ded30a7e2eb Binary files /dev/null and b/models/ppo1-icm500/0/default/manifest.ocdbt differ diff --git a/models/ppo1-icm500/0/default/ocdbt.process_0/d/18763adcffc7b993ea4faafcfc8e48f4 b/models/ppo1-icm500/0/default/ocdbt.process_0/d/18763adcffc7b993ea4faafcfc8e48f4 new file mode 100644 index 0000000000000000000000000000000000000000..e6af2962f5caba1e866287b1d7cc9482791e9bef Binary files /dev/null and b/models/ppo1-icm500/0/default/ocdbt.process_0/d/18763adcffc7b993ea4faafcfc8e48f4 differ diff --git a/models/ppo1-icm500/0/default/ocdbt.process_0/d/4e996b880df24475f131ce5b73da7f6c b/models/ppo1-icm500/0/default/ocdbt.process_0/d/4e996b880df24475f131ce5b73da7f6c new file mode 100644 index 0000000000000000000000000000000000000000..970e98d09ac07bb679ed747d8bea6c27ac253f9f Binary files /dev/null and b/models/ppo1-icm500/0/default/ocdbt.process_0/d/4e996b880df24475f131ce5b73da7f6c differ diff --git a/models/ppo1-icm500/0/default/ocdbt.process_0/d/5e2e6028207e4d2739be514a4100684c b/models/ppo1-icm500/0/default/ocdbt.process_0/d/5e2e6028207e4d2739be514a4100684c new file mode 100644 index 0000000000000000000000000000000000000000..4abdc353623982deb86a7fcbe7d23692a3935dbc --- /dev/null +++ b/models/ppo1-icm500/0/default/ocdbt.process_0/d/5e2e6028207e4d2739be514a4100684c @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3d7b78ab14c96d38020560ba0f56638e9526b8e4a3b1fc46fa605b4ec897b130 +size 526053376 diff --git a/models/ppo1-icm500/0/default/ocdbt.process_0/d/7283cfb568f3403a47b5f10874196fbc b/models/ppo1-icm500/0/default/ocdbt.process_0/d/7283cfb568f3403a47b5f10874196fbc new file mode 100644 index 0000000000000000000000000000000000000000..4b1a203a494f8d5cd524894f631d64fffb91355b Binary files /dev/null and b/models/ppo1-icm500/0/default/ocdbt.process_0/d/7283cfb568f3403a47b5f10874196fbc differ diff --git a/models/ppo1-icm500/0/default/ocdbt.process_0/d/de1203851f6380ed6e4899697f68e66a b/models/ppo1-icm500/0/default/ocdbt.process_0/d/de1203851f6380ed6e4899697f68e66a new file mode 100644 index 0000000000000000000000000000000000000000..85b75d74264b2b77d405db7dfec0e2416d9ddea5 --- /dev/null +++ b/models/ppo1-icm500/0/default/ocdbt.process_0/d/de1203851f6380ed6e4899697f68e66a @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:966ca3d2af1f85898766c03581e59c3f8dcc47366007bc222d8ba9d6b6313196 +size 53161984 diff --git a/models/ppo1-icm500/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-icm500/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..43e8ef085ed914dc1c4e2efaf103b30da7ad6764 Binary files /dev/null and b/models/ppo1-icm500/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-icm500/checkpoint_manifest.json b/models/ppo1-icm500/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..4fa149b8b855a4aa4849a1023228f32bd99a888b --- /dev/null +++ b/models/ppo1-icm500/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 29000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "1c4993402e1c073499a1d4370898b3e8e9a316cbda4a5f0d43e266050487037b", + "runtime_config_sha256": "1c4993402e1c073499a1d4370898b3e8e9a316cbda4a5f0d43e266050487037b", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "1c4993402e1c073499a1d4370898b3e8e9a316cbda4a5f0d43e266050487037b", + "schema_version": 1, + "training_contract_sha256": "874db0402e4868017b1ec5039889f0998b718b4f13b4d7dd2316ec29c627dc09" + }, + "sha256": "fce885b9545238543d4424a21311998d3f78e81f5ee2cb1d46a59f065c2db90a" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "ce19068792336547da5c5074cc64534d59427631f8abcf89d912db3e76452f42", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "f7a9055548509af01a076b33642c68f430b473d38a2a1754314b6d96e6a474f5", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "462b5430b1c18a1455c061f7d6eb2f19757972978814c66c945ad05d262849e0", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "27df75611f6460c2e807318ccf6200a94a8314c3b259449a9b726e51bf85d4e9", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "6071b5422e696c6e630a8ddf3996ddc7c0a651c081be55557312c7cfec7f6149", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "c046a1934ace88fed1e8bd97a2992cd19c48b29a4bdeaa631ddd6d8aea939a7b", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "172479937faa1931a85302e0865948035bec580fad9367ea913e4bac282ea1d6", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "c23e8f2728fd93feba87ea20bd498ec3f469e95a22c778daddb825b7a955a8e2", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "8e43791a88dc288fe4116c654074e2e6b50305b04cede32d27d2912c727c0af8", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "ed8200e8dbc12319ef6977bad8064fa6ffd6fb8861429445adec0d50f0276ad1", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "db3db48c7134205da64cb52a2a655fa35c04c3b0eb8d28d8fad1c246796dba84", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "ec7766ce8f792e55dd67cc4ded19e7661c0a0bcf491c4d3536b2f90d447cdc8b", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "bd467400c7aa3f680daca2eb31874199dd9ad2f6cc08ce19ff7fd4a733a1a5c7", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "e57ac101750da896182a1bdf05dfd39e9590c18d0e93126f8b43c30ce2a2e389", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "673a8a9978c74fc49939cd2cf34df2c73536115845038a58edbc92586c795500", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "ce7e8ba57f57f21587a71aa4178de797b2387782cbc523da1de2aaa9a8daba6a", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "2b2dfbb153d127eb9a87868c9a97039f2f0cc9693b889aa8a2782ecf9e0bcff1", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "65e390498df0013f567bf34f03656e03c6606ac8d09bdc14396e7ee9878fbc1d", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "739132fc2863621925b0c42186584ef393719d3fe47c46c54e4bb953190efaf3", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "22c3d98cb4009c9379111b6be2c66ca54372c2b81574b9bda7f5297f68bccccb", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "51c38f1c5e20adb2d70b7dcbd9c53d06c13bad7a1544ed6a9870f4311c8b622f", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "858e6e4a994112297619081279629e6c16cd8af27c6158c8dadba06b6a5f5157", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "78f70c7dcf6649f4a3b2884b4b0e63a16868a347ebc09584f748dcce79df2427", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "f15293bb5a7b26622a221a5f7151012b99325747e28685cf1e71dbc5e494056f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "048b4e7a2afad3dc39dd9f089f22e40ccf470463b8432d611facf7b8845980fd", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c046eadad8cf0bec1a3a477db35d56f7389a2e78fa4f3c0ca2e11199e26e1925", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "4d7308b5f391f8dc7c6f90624ac991664c4848aeef3f7fad01cc21cc70458798", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "9c2d8cc6253b490e655bec7e33f76c6a3401d9024db9e3e237980aa40d229f3c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "b50730e10a360eb43b2a1633a3fc75ea68ca0d53b2e257f82be05cf3c2767515", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "95e118f8d705758c68f4765f4816c1a74197f4a3787a7f825f0c5d0488c9f84c", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "564ee6ddf899f626490d64257a6a0c564280a83d1648f4c7270a4daca84043a6", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7c3b08f503cc8f1592cc07342d7c74028a6408bbd4b6cf04cc26be4a2e161dbc", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "5a0ed72cf2babaed24e1905dafa52cab2d30ef5e8cf9eadebdf780129a923065", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "fc2d552ffda2fe8aea572f1e712f783bb32b992fda695bbf0a86c416a9ea7ffd", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "13ff28d046de27685e8578ddc4a352e6701914e50b98216a7948622885479d07", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "d065de2e3e130494a9f935efe6116d1bf636c6ac9fd3a59ad93af0ffb7a996b1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "94c132dc4eed5e92f4ee7a85e62e186631c136148dca0dc92cfb7877cf07da8f", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "1743c7ec69775c7bdd611a2863aa60f426605b454571bb72e7bd2b62ab1a968b", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "ce9c897922a8ee24f32e7c4f12fbc42689a1d1e9b0e0050789d8e2d0d96dedc4", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "14921ab06bdf814cfe514f8eed5204320bc79a592d9c8a2b96e6bfae21e7a85b", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "bc7c42556c102e487cf8def1cc303a16ec37d1a0a2c60f364b9a7e84612a41cf", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b133edbcab775880894537a004d3ca03b5d49908410c54fd255c8c20c04992e5", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "4e3a6b011eba69576b527ba9da2a2abf83dffc80e439f98117b0ac91f8bdfeb1", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "89d2ebed762d77097e9f80103c4a12ae15530f0574a0641422588251023aa463", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "b44ae97174dd4cf3883de59d6fa27575b0c12bf847813477aa082fe8165e9e65", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c20af0ebb1f4d46b61bc4bc19ece33916c298900299cb21db66d45e364a013ba", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "cbf47f1ab671ff787c9953a04f5939a064f14300e9f5c3cee2a2798f3c7c1a5f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e7a5bfcae143beea3edcd470d7fc76ed48a21ba8ebbc9e833a53736fd7868917", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "337b81b41941e62c4d476538cea3d61237b1509bf532ef15d93555a0e28bdb2f", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b1799f36421660eef6825970f88f4b00639f2faf006ae716b6ce71e0e51cd47c", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "15a38944dae4c9b0c6b61cfd94896a952d9c41c4ff311e88eca4315b3d15038d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "d354b00c014aeb82273bbbdf4c378162a9ec57d300a3728ee88af1ba0e5e6f04", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "428e045586dc4232db498040dd5ea8d892ad1e6fc340655d478369001133b8bd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "97104da4abe528a212ec27026cbbe804d9b1d2c285f3f733f800b1c29dde0d1b", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "795cadc1720b864547423e475d0ff1848c22b0a97c3a7d3464db7a1f3085e513", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "fc84a9e28cc64b2dabdaa20e38729d4645eac2cace7fd23e8ce5cdc67da9ac89", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "7250fc26e1405b0baced445bc83ee2e71b991dca8a716fe6458b4a7f3071b2ea", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "abf42a515c45821cdb3e36848b8e4cc11fe2eeb59227db81dcdec205232afcf2", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "19b6eb549440186d9d5bbd2f4904b36c223ea775ce2b0ada39a3e99cfc3cccfa", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "f9ab1584989160e0fc4a16586fb6f6361b14bbb3d1647606f6fce2c30f8ece9b", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "17f10841a33bccf514ecd44acfbbf84bdaceb769b65edbe841e33d23bcd79c49", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "811378503b1deb0eb96d2cbb6378968d90aed20317fee8b7fcc07e599fb96b23", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f45660ccbc034bce0ad0cfb123f3abfb89be6e2ed2dddc02fd58d81097af0bc9", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "04580c428b3a387f4c0ce2bc4bca7101f812457eab2876cee383fe5575ff4819", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "e873c1a023871c073341ec0f74aaa9b52cbaffae0143cf8f010520d8a2d62428", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "b0da05868aefc73fc8b48271271d726ea879a6ce73c6c6da26f2d0224126f6e3", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "ec3ca13c3dca33036df4458e1e392935b6d741260acba5613e7fa714dd1574c9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "47b8477aae2176d96d4d87c6c5968dd3abb75f1873b60912fdf91e5cb9aa0673", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c232d624b90acb9e6c92d2cae0846fe931febc59832c7c1e5c4e06758db17ae8", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "f5c8253652d6b3ec012219802ccd0e73afa20db7d843af8efeead99c005d834e", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "824ae5bea7b9461749c1e6d0fe293d57edac0dcd520fad4af5dd3a04dd5e06b2", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "60fbc90eae9f491ac9e258abea1a900caeb4de7398f7e9d3afa6acd2f94ed72f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "9345aff27746ff287f87bed9d377aedc41c5e366c4f447396039d049fcac4949", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "9c2bae04ceb305a70d2a467a00593e93248a29ada602f4ad0f5c7a9b3d434366", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ba9edd8d4ef1ecea64f572a39feb9108b18178d43d92662bed14bee8a6f31724", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "33e0b579ebab59314cb50573babf7ee5fb98887a4e06f8631ad7588f0a3f6303", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "735dc53b07da65d4e938261b1a29d6bde63eee2a9be2028437ceaaf173f29491", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "441f0929557c038dfc9e797be0d7bae7785783a660aa3fabbb2b65acb813f606", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "ae606390d8ab407c7d5a9d0bd550d4f532450287daf934b90681d417f6e92dbb", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "ac14b1de114abda4231b617762966ef515c5c368cc62c95b15841961a17c200b", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "3d136186426e673efbe7dcf3a10f70a8325789149b4913bf9f27f04bf216e10b", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "e6136b6035429abd4ffe59ad7c9846639d432829f58131eb9fb97562eec56c42", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "7dad1e715801dbd6145ae61282b984a1a5e95807e5c0aab2467e2a26176e0aa7", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c94515b431a90e9730db6a0ff1d34bdacc9f3b272b8bfc87633ef526210531f2", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "1d81678eae3577c6d99483acd5eff6f9dc3dfe74f62060f0e479bc1cdac9cc1f", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "03127772b308293a35f201b4522c361f5a7ae9cd1576bc456aedaac4db3a6bff", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "66102b0065797caa4bdcd4e79149294f03594cfdc445f2c8c143fb9d3faaf5f7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "6ca50138cb76be218036e15ca3cee9e9f00d1bd91d7e86e31c1af0a4f7ba48d4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d4c6db96e0b491581479c6d0854f8d1687cb8622263e7d08703825b2a59c4821", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ddda35f276f99c32514f992ef89e9da29e8686b1c63ff2d5dafdd3639766aac1", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "7f20382ba5706baf4791996f228b1bdc7891d8d54849337b2e2f4b98a522c213", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "411dd100fe001ec59e521fe8d1365e37702b68011319a1aeeb53eb6cd3cbf0ca", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "75bf8a848e6f14b06181fd12ef06f7da2586a6e44ed8561994bc139803eb344a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "8499a6398109b91dcf151f52a0427f3189c41863aaf9125d8a3051903c4762eb", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "0f39daeebe98607b9eb81a17a6e5ff45b01d4f6b66ec56ca26c06b22282495b1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "2581e8812f8036dec83a7ddcc4c266aecbef64b49b9e22c661809786695129c5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "5fdd73f5625aa21b7f10bf1539c69122b09759713fcbf328c9487ab16d15431a", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "e7ea12759cd611b64fa7430607c265ea44f2b6e87912f62fe49a436f56839166", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "f41fe63e0233c143e46760ab94fa141b768d5dfa5c51fbc2448ac165ba544f84", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "6780aed075208b695fadb8fbba5173cb3de5a918be3fcc6dafee249a88068c71", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "b1282f555dcc8af3773ee7d9ba562a8320e31364ef0dc3f80b251f12513a84a8", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "0efc4d0f3f80da87b9df2d083aa9cdb14c0eb9bb291407a119302cbb43c1eeb9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "d1947620ff7fd3cc691a9a297bf10b53fc87e9fe792d104899e872afeaaf6a49", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "86e9eb42fa05e323d672cb7f353a44c5a5586cb0b355af51316e71e9cf01556a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "c4e54f246d9d33a0da10fbfaf82bdcf2cec56e4f024ffad97bed0566e91e4c8e", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "cf288861339b6d64be1372e598dfc8629bd1c4f902e836eebbe0098cd770ccec", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "cea82b9f8b8813494d0dcc7b9792d51c7b2ecc587f2719b15284c78676cd4bbc", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "419ea2336b3230c6655969c412d635a115c2dbde955e433994488ad4e9506a8c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ed9f1f3cebabddd2e2b8f2e18fc28104097bc03d4d285396bd9d5623eb084bdc", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "dac2a33ccedd8b35c67113b00f8b9810c0475301b4ec0360b201685929e2772d", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "24b069ffc5d41965ff81a3c745c0a41970a4c98b37b68890ab148ef8bd9f2245", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "6956b3e8b5cfb4390a4b2e85c60a2344e42fc8b24f3c1e7229d1e15a758a5758", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "f73ce109e832d4176b3a74bdca9b79994fc1b4595f5692e47784b0f39c28eab0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "feb29cfcd6eff3050de1f78cf28a9fb0a3b20e23f39ac919500eeb33671dcd9a", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "09ce8622873c1c4aef5d6af3e81651579b8475c04f4a9cd0847b0ac2e2c90519", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "5bea8ffb42c368da0c2b054502c3d9ac2034d3b1f5af98f25104b976cf8ba801", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "d891e6ff9a8a4881e0dd678c5931711fdc8b440e8bcb6a8f156d8ff60baa9c26", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "da6495eeef62e976aeff33bfce6abad62630103a085216e9ab842989e742dde9", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "e32e6ec1c25e61d59b724abeb9e09efd896462a0c299cf4f9f0a9b537f6e7408", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "3e666b9c24836f7a96c823c3f95698ddb0247283c7d1a9caae81f31f44783be8", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "afd272474fff184dbcbef4cc984491c4534fd6a2e9f8dbd08a2e286c3345bfc1", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b8f29b60622ce12fc350241735bb5a5192f9aa637327495e4b10cd1d79123bfc", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "81a459506b8ba66224f8ca1caedde443d4ec2715b66a7484efdddf3090bdc2fa", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c7840880d3bcbeffa2c6cd0d5a42fcae76d9b99ba0bc6ee5d4c60d09ec8539c8", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "aaba8d53d30ca5d7b7ecf926bc8309c8688673a6973dd2ad7eeee7b3c6ef11a5", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "87c2a913dc5d3c87466045d9c53b14dbe6aa694dbbd8699a5591fe00724a373d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "56f871f5c321dd49d357cad82ed87be639733f3e38f4e2c6da638b87a91bea85", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "dbe162943b417a93d9fa153710a872c3dac679d0ed2b008ad17731b215c58067", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "f88f1d8f49f9389df4b905d8db087652a81d68ce821e80347992de6a7053c53c", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "3c03287eee54234a6cfd838f55858986ef7bb461e5d6cb8f647cab73e2deedbe", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "05becccbc505b70f7f86ba2e93ba35efbfb16df7846816f4b386fd785e521c1c", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b1fb3bb7f8e68d23150a4290c911b6d5515a64d2559df8e2c51de887b00dfcbf", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "a97d1e6b7a363ac3364316a9df5978d7ec0a5f161a8bfa1d326e4187962eedd2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "9340579dfc7c492c597756b76cec857cad4ff40899cd47d899d79bdad3cdbd9e", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "5a65e084e93f8c41b35e6425f3ab9158868cfd179c99016fbdcd24eb137f27b2", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "713de1f64ebbb9c902f42ebb9d05f107771764f4ac05e617da5c6c6ff674d564", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "730f62615049d0fc2a5067b986aceb5a8cf37c2730409b5eaf3f939a7f4b370c", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "3db7688dcc2918a72acf939af65047b670db98a1dddd12493bfeecf6175fbd0e", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "81c4bee7548ae2fac72988345fe460ccec8ac9f4a5f312fc34736477cdfceaf0", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "d032cc7e61ff53793f6f040b49d80ece5a3db8337555f8ac56ab66eeab58bb50", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "71f2766eb63dd09d788697ffaf5102dacf6fe61c913d2938aa9632312acae4c8", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "edd6a305e1a43f6b24b5cd747ad454b5a7045a4282d8542904db6d282e8e50e0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "3b7fb1e7de17bda522b09aba103ee0b80d9460bd916b4e5f9a01bbd43c4000e4", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "20a84eaa31c88dadf3de4a1e88bfca98e574f73e88b384fb75fc3a0099ee6937", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "6c15f99d82a6df9a4836ba8296e46e84e116f7fc10eca237c4242fe94a256d16", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "584eb1b2049cf30856c780995f96452d0b629dca2eb08a1696a72187c576170e", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "8e7fa86df2936d236d042bb3dabd78e715205439a6ba671bfb6326a8c82923b8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "a3600df77705a53eb3731e1a72d73c42738a9e8c62dfcbfca57a26b517c28ffa", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "3a8484a80b9e9b913e6a43ee8b9b93e467a97256ebc46e4e4ed526c1b30da848", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "cab289456a3b1ad0b47f25b30205bc8ed983e0cb092ff88a0db63a5f413c76fe", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "efedd5e963ff41ec16e93c390825b6be846b7b196366fb5b3b41795b1c9b103b", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "8f7ca0a522613fea8aaa1aafa64bc729c950ff385448c71c2f92cd7af2bc29b0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "a0ae60a426ba0df0b2535f6d049935ae4cdc1b81a11108fc603d00106bb02bf2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "05c4b2165c2a92cf7e8d82520c902b1486e96e19c0bb6d700a708b9159f36657", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "62aa3c8c7f10c083cf3e33dc13e67e0785a3eafbd93f1bd86db5b18a84dfd344", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "084ce5f512927abf09ad5f3c21485c38ad31abe5b0d94033f51d5811427836d6", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "1696fb70850ea5fdcb56cf74193c9ef3eee04eed7920530c81336dc1cb72a54c", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "82e7481bdea5b7a1a5a0c5fe68bc8d8c3dc4549280ddab1c6aa1c73532c62fc5", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "ca874e628fc1bde67f2c5829106074d9448002b7b3e898e8a861ef42a73cc8d8", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "019d7f09054954a5a305ab9cec6e89ca3b9e38d971a4a4af99f4369950dd1478", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "624ebcc96eb064b13f076a3abf919ff0562cf44d0c63cb6f4e964f0005751042", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "71a23618b03ea31bdb613305893518490ebc6469c5b3618452f6edc6dcc83ba6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "b8c4cbdf8b67e84020f8637959f1242c6627134fad8a1990d64f66499d417512", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "88d01cee048ad0847268c3b60dc1ae48f753d75f7441a8be83d911cc932924d2", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "29e112190c12975e5740620d1ab62c3571603f0fe68c0c28f8c726aeaded75f8", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e834408e9bc1c68e7e00303adc3e6ba9c826d8d00046e7c53b7d2d3b6dda61fd", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "58af195520045573841856c3b145bac7ef1ab4320b7d4859c1f31e940327e3c3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "3ccd59bf807f5e9d1a8095e2f30f8b5b577748a4ba5fb5e1ccc34089fd80ee9f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "2e506a12e477a7436414a2e25f1188d58ee5fae421981c5cdbd888e8aa7a6d60", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d97b280af8fb9c94630aefa72e4833edeee4ff786fc0078478dd75c455cb82dd", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "bc8e0fcee59c3ff022f14cfd00baadb227abafc130fe8cc60f37f3b0e74af6f7", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "85469adea27f8bff2ddc0c6d6cbf72ce0c2060e7df6aa14f83cd5084c0766dbc", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "3fc6d39b9d8b801eedb99d439b950bc5e8151ef64ff8f94b12e5605c294fb58e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e07c4cd6052941c8b6f326e1de81308beb91dd50eed4f305ea281c92e130f9f4", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "b40668e15be7521b3dace6e04e5120c7abeb3bea05c6a0734abe0d73b948285a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "4a2105875b7d9dd07598d611acd71f09bb2f1d28c367a208134cccac964a0812", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "32890c83149ee7d2154820487a43788f8824329f1f795a901ab38cf815692836", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "b5b7eb059ff3df277449e6e8c89aa6f8822e73da25840d688d8cd3085e3d2465", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "3bcc57b6eec20a4779b5678e116ad387eb33b372afe8e00079710e4cd673c834", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "4cd4591b5596c4b7beabfbc328cad6d8a8e183ad7ea2cf9e412a2187c9b7fc5b", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "2c7ad8056865efe26834a65c73361f9f047d162147c541e5e99aba071e6160d1", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d34744c11739621dd69618f9c0c53e4ad17795a6ed1a986a36c03816f34cf69c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "0a03df31d65064e11d973ee43d3a0c2cc3db161e9af20b707c2e7eb2610d16a1", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "6872f6730a86b7ad9980f4d5da2cc49ff8e3f07a9a766130f9386ec94882f8a2", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5688ede7d5ac35a75fa50d932d586230a80407a3d5c3fb1fdac19eb8ac994818", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "3f33b7f29704c617def1ee129ea2ff8fcade1cd878a716593e1ee450ac8a1a2b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "c7700cf04f1ad51ed760e7120b6eaa47fe6f99831042b9433ecf356f44f8c900", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "959802398a26be97bb79f49cffcec8b152335b93666a09331319c7166c89be73", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d22da378dd0e53c4a400c5175c8a03c3dc044acb388d663968830faec8d87ef1", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ffc1260944dafe5847e1e8d3a49f2897b90c2879f764fd52a6295ec00cc82c24", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "fcc71b280e9112270daab3c2bb58befb437675342e20d6ed5ec30e5e3005e13e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "76e6ada4b43c604ae13c68c832e6c051b1bf76ca64326a677bf5654cbb3d3028", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e99476c15b8a71dfd7d0834842549f96d5df9ebed7af8082e4a085af379b50ec", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "c03293f102ddc049b193dd7ad91b85545fd8ef6d6d3d36af8a6950e36dd15470", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "0a43b285d6a8d7a1e222e0c74140e5245707545bf9ac93976485d0fd3be8f891", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "5e106af55d7be36f87d007ce73f68cfd2a077295fccbe9f581d0882b78fda110", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "839e26a1a134b8e13ec133a4a41b6c64b474d36b3155d283e8e0803448fa1089", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "b55edd6708c77918b0030f216b36f45de5f80bd91100da448c61b17417b673df", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "671ee538f355907d4505035b9ed6d3618697dcd85e4256c446ab832bfd8d9501", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "1094a930aa29e83f610d0c41c84fa2d8c89da134ec47f75c3f57447c32fd1bad", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "2b8f63c6ad16a0095124cf4051e5e37bb873201873908c1ee354bb4b4b2e30cb", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "bc5e0e681868bac93986fd5f2d6a017d232b095233b886a8180862fff1fe753e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "9ba4cda1cb633e1d470b1942a091677390dedb45bb00e27dac04e780fa04b030", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "515f282e9193f59cbfed5eeb3e2856867262d56c53e1d63fabfdfd44445fdbe7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "0d7d88ee26bd41363b6f870980dbaa4d79c685f1be855a9bff2b19fdb674a71d", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "dd5dbac90411615bfae6ef27e80f4633ef65e6a53c8859a638512071e0b71654", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "64d6401e39d5d6a314e3ab4e9f0ab3d993576ed3475ce0df08b0e32aeccb2492", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "fced082f352d8584d1dd0f499f1156116bb87dc72420a7e4435d3c9262b390ed", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c3ac1c6c6c225ed89da2c6a8bfd33f3b1d9bbacda8320545615520e485d2a9af", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "0c594249e41afb9232ade0e8281d88db4b16f0920bfee1244318de6618c62fcd", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "2908d1321ec6fb193d1916910e1eb81274de2dd6a49e577fd9460e1650668b08", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "c5467557580c9fe3897ff4db77c0ab7b77454e42d89267836937b3c02822f683", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "35ccfe122edaaac97e253402a53ddcffb4149bc1de58d1a12946e24a92355470", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "c32b86733dedeae9ce3fa279ff86ac94f48739b2f272a67fd2512c5aa741ddcd", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "59f65d808de25ce48a6ce885482159c01acc6c9e4f6d7783723541dd0de68ee4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "6ece740ade6d7ca50ce5bccefe64594dc653a1b5224bc1b676eeb1d8809ff2ff", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "c4b13708bbe1fa1f80a00e636b37324d80812157d1d0305fe5bb822f09db05c6", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "8a96589fcfee755adff535cc655450b41ed8f20dcedfe907716592612eea8be8", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "ca3d9ab623ddfc927134517799c2c945279c69675688660d905ff20a37b0601a", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "ec36006753789f198a827af5564eb7e532a0ddd8954809ce15e0a6bd29b2ca75", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "a8272bc1077656800aa59435ff738cc6142fa4230415eef4ac4caab5c72bca0a", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "94f703c3fcf4beaa472979ffaa47ee76d870f2eed5aebbbe861903ef1770841e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "2d944d136ac429de59f870d1811facb624095becc39c19a5998c8540ee7e8969", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "6888306abd0df52b8a7fe51b45d5d60bb961ba3beb46343bccab83a179138894", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "8e5e3bc42314ee1309c83adec2df3bcb8023e86a949f8ecf85b8a09406a0901d", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "69ec7bc7f7161f23e4c268510f79b7dce18dc1411e566a2da26927f7dc8c77c4", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "dd6fe2b202a2fac91b2c49af5901bfd243412fe71a8baf1736bb079ae9498307", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "f210b9c2fc59d4dd6e2bcf9d75b41e3db322847c80986dbc6fbea1ac3f590f81", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d4130915529891d4e3f15d2f794dd0b87a580bf0af24f928dc9f63eb92c3295a", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d862520873b465f3acdd1262702f921f33b954736acf01b7d302ee18d6f3de92", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "5c59a3f451914cff884320573d99c0ba558336f7a5e7550968762ab9dbb34e22", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "58fc83cd0bd996447a05733bd03d6ddb830b7deea4f5aebedd78eee160b8a4e9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "1cbf165c72f7c3d2750c8950d10a037fb9ff9972e313988ff1336bd7c595bff4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "ac6575bda7a1ae3f0959dc97a23ec5d0267ef8134496ab0ec731cd9f5440d2f8", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "168ff24195b47b5cc7b30288265c05b88115286a5e982ed4f05da31b63e56a4e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "15533d32ad14b75bb186c5020786118acf17e3fbb87ee2cead4f988602add74f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "16ccdd4cbf2079d0b04de9feccd884a138f0b0429d706edda4e7ed76dba2c8c0", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "fa3f178a753f4d4a955a730db3c0ece183840874119e302147f17252291bb733", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "09a4320349c2fdda87fcd70702919c962e72ac3227a6a9bb6c4d07504a4fa08a", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "31790e98c19f2281668d3c6ba2d5fbfb9fc5a39597f4aac9c131b5aa740bd821", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "887b07eb76fb7e5b1c7956da3dc3f086a0b2bbac9310a6f5205992a702c3eb19", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "ada6c2df5371709c90471687583222d09bf18f2b8f821c7a729c5053f0fcf51c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "27cf003a4b4a93770e5f928d113ca9cd689ef173502175a8d817453d4fabfed4", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "2871a6cdaa1d038c6a14c9e6e2368b288744fce3b3822edbdb2efbbc6aa2f20a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "f960f7d3025c120d9ef56081457738dcd5fff94bb1c4e9a1eb1c90bf0324e382", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "af41beeafd9f3b1994e20974c4e7f0132814385a2a581c5d19f26973c97ac17d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "901490583c912991803d344ce229c4fa5169468e3674b8655589fa123ef90980", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "bd42cbde313795978c582332864856dc062d769b66a5caf511c3917e5d526f04", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "eaaa2544aa7e42c9be17e6afbe7e83800ee9b070dd8ad70a0b169906a3b29a8e", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "accb18b3801a7f66355508a64904d4c5d9ee0539913ea754c6c163d71f096ee5", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "88f030c7832e5648ba93b849e100d2ef434d67a1a1269a37400d552fb0378984", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "c06f35cb52c4ed915199c45abe44e004e42e9e1535b5c8ef81b93e26c3f40ee5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "8bf38cae34cb08bbe572faadebb967d9752333e00b2bf0e06297a142aed678bb", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "cbddcbc13030e8e7e9d15ec7e6bc3c7c697e5b62256b9820009361fc6d7b97ac", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "782735d8ea7c89aee8804704f0fa67e0023f439387cd4176d4a75064b11bb63a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "fb31818b37c9ab1bd8790ca36ed3346ca3e4f20acd094593010a0a0ebb7362f0", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "e9de104584a344c8dd011bd2c3c267ab185a9c40bc1e46d6036d4d791367b948", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "9873b1fdafef2e8652ec2eacb7515e9a897b50341a812ad6da34cccda6038540" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "0ac13e8dfa812f09151d2efde6ad3904466e1c46682f550a81ec629a1d91300f", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 29000, + "training_contract_sha256": "874db0402e4868017b1ec5039889f0998b718b4f13b4d7dd2316ec29c627dc09" + } +} diff --git a/models/ppo1-metra500/0/_CHECKPOINT_METADATA b/models/ppo1-metra500/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ca0b3a73ccbe15d22db5759397b2ee22b3a03564 --- /dev/null +++ b/models/ppo1-metra500/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786225530709294247, "commit_timestamp_nsecs": 1786225531360423368, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-metra500/0/default/_METADATA b/models/ppo1-metra500/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-metra500/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-metra500/0/default/_sharding b/models/ppo1-metra500/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-metra500/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-metra500/0/default/array_metadatas/process_0 b/models/ppo1-metra500/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-metra500/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-metra500/0/default/d/fd19d4383cd42974467aeb5bf1554a4c b/models/ppo1-metra500/0/default/d/fd19d4383cd42974467aeb5bf1554a4c new file mode 100644 index 0000000000000000000000000000000000000000..636d3359ac105e113cceb630786f86c35d58bbba --- /dev/null +++ b/models/ppo1-metra500/0/default/d/fd19d4383cd42974467aeb5bf1554a4c @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a54f5f9c05ddb5b9e2b9e91f6f96e6f4bed5b503c8ddc8212ecfef0bf31ef83 +size 112860 diff --git a/models/ppo1-metra500/0/default/manifest.ocdbt b/models/ppo1-metra500/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..8fd9dab655040847b22b4a7bdb81fcb1dc27fd4c Binary files /dev/null and b/models/ppo1-metra500/0/default/manifest.ocdbt differ diff --git a/models/ppo1-metra500/0/default/ocdbt.process_0/d/1d9d018ec71450f96ca3628841948c3d b/models/ppo1-metra500/0/default/ocdbt.process_0/d/1d9d018ec71450f96ca3628841948c3d new file mode 100644 index 0000000000000000000000000000000000000000..bcd73e754d2da9c40a3386ceadbc5efc43deac9e --- /dev/null +++ b/models/ppo1-metra500/0/default/ocdbt.process_0/d/1d9d018ec71450f96ca3628841948c3d @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ac4f76f6a56f58f938ec21f54e4ba3ea9f8b7499145ca3af39299b3c935c8164 +size 511373312 diff --git a/models/ppo1-metra500/0/default/ocdbt.process_0/d/79b3d14665e9343193c53e8fb237c828 b/models/ppo1-metra500/0/default/ocdbt.process_0/d/79b3d14665e9343193c53e8fb237c828 new file mode 100644 index 0000000000000000000000000000000000000000..4b1a203a494f8d5cd524894f631d64fffb91355b Binary files /dev/null and b/models/ppo1-metra500/0/default/ocdbt.process_0/d/79b3d14665e9343193c53e8fb237c828 differ diff --git a/models/ppo1-metra500/0/default/ocdbt.process_0/d/c4cd899b11e2b55f75cad9e4477e26a5 b/models/ppo1-metra500/0/default/ocdbt.process_0/d/c4cd899b11e2b55f75cad9e4477e26a5 new file mode 100644 index 0000000000000000000000000000000000000000..fc548d6b737fdb6e1b76775a50fb772db88b511e Binary files /dev/null and b/models/ppo1-metra500/0/default/ocdbt.process_0/d/c4cd899b11e2b55f75cad9e4477e26a5 differ diff --git a/models/ppo1-metra500/0/default/ocdbt.process_0/d/ec10c9c8cf0eb47449729e032e204e75 b/models/ppo1-metra500/0/default/ocdbt.process_0/d/ec10c9c8cf0eb47449729e032e204e75 new file mode 100644 index 0000000000000000000000000000000000000000..042e4ce045ed0b5286307326f6a0d970818abd65 --- /dev/null +++ b/models/ppo1-metra500/0/default/ocdbt.process_0/d/ec10c9c8cf0eb47449729e032e204e75 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a4e510ce0ac6d5abb66312b57493e5c9470e33b8e10d4d7717c9a2c597cb4889 +size 67854336 diff --git a/models/ppo1-metra500/0/default/ocdbt.process_0/d/f03e4e2e123e6f9d57c6f3ab1559f880 b/models/ppo1-metra500/0/default/ocdbt.process_0/d/f03e4e2e123e6f9d57c6f3ab1559f880 new file mode 100644 index 0000000000000000000000000000000000000000..970e98d09ac07bb679ed747d8bea6c27ac253f9f Binary files /dev/null and b/models/ppo1-metra500/0/default/ocdbt.process_0/d/f03e4e2e123e6f9d57c6f3ab1559f880 differ diff --git a/models/ppo1-metra500/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-metra500/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..8be384072137bdca25fcbfb540c74dca76c3534c Binary files /dev/null and b/models/ppo1-metra500/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-metra500/checkpoint_manifest.json b/models/ppo1-metra500/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..637b3a158cd940d3de29e45f7b4b72c22e711434 --- /dev/null +++ b/models/ppo1-metra500/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 26000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "c5174e62e7fc0018e24c7334b4ba09733f52dd00b8da5d0ae2400a748945d964", + "runtime_config_sha256": "c5174e62e7fc0018e24c7334b4ba09733f52dd00b8da5d0ae2400a748945d964", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "c5174e62e7fc0018e24c7334b4ba09733f52dd00b8da5d0ae2400a748945d964", + "schema_version": 1, + "training_contract_sha256": "e55f0cd01870441f33c9f35544428cbaf029edc34fa1637973d29f64a5242966" + }, + "sha256": "4964623fad275528492c26d12d491daacd27c207d0515ed0480b4ff45007b9cf" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "e2dbe455bddbd1ccb9d7a639dcb142524c5574fe05bf999d75b9de46bd2357ad", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "3c3f098382946bf6dd1e95ed85142707d4f7c56cb4ff612a81b04306fea4da73", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "881012ccf458954963e3525f5cfbb9856584c34fe9ed0750702dda2df4455d5f", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "7e490f5d84f4fe5938f17a5a4e6b791d38345632da5029e976ccc4f42fde828c", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "a6dc829e1dd288d430f2a03162213c9c07c61b81afbfed51a31ab4a9a0173115", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "6f7952a4564a7966204fa6cc5f05f8a58f16a44019fe10d4f90e1cae4b7767cc", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "3b30995cea694aaec54993e4942d18c3bfb5b5d187616e2cc3a41e03d225ba9a", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "5439083b29f3c82e4e52d470a0ab87da76baf3e9ebecab0fcb457076918b8a8c", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "2b4a35cb1064df59fc6903ca48b9b69b6bc65ed9552ab67bb6b3c9fbe6dfb113", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "9f6a3370ab45cd7dff00f9eaa4bd67547728559d688297e0d9e78f542a4420cc", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "775cb8ff29ff242efabad8fa710f39d5bfc33152c22029eb60aa16b423af0284", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "05531147e0c42c72e910083c33824d55e29174d85eb4ba073cfee4d5a07596f6", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "c1d2b20ccdf323266fbc0190e50c4a829b5662c0dd802a972906f557344cccc8", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "e25fb1b297960027938daf2fb133452315affd5d28cb3e1b476eaada7e28a32e", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "ccfadb6adab0a5590efe470d6f2c39ec1c59b81ce1e6ff04b1e14b9e0fad8d7b", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "1370f60e8cc0723571803401824b87c4c67a6eccbb91179c972926e81a69c7ff", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "4bcd135021f1be06e0f5f3ee49b434d941aa6a272e046ef8af03fca87b5d9ae9", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "dfcb5c7c878ec97526e964e3b60e00c7aa6524b56c8311794d57b8c8f145ddaf", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "846f01012972f4c8321e3617ecaba7d556f43402640b89a8c81799095487aae7", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "15fabaf933a8144e2c8c54a5a592930dfbb1f0a4246e809175f2c6d4241f9e2b", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "061301970416d296eb04dce940c853814d939cefaee1bab54ba136bde44b5c3e", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f4f7eec3d65bf26816bbf613f0ae1dfebd749ac2c9cdbdb6e798704902ddebb9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "1f2656b09da69291ec9b224ff07dc3953950dff726cbd92ab3fc5061d6d1211a", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c9d9d7eee44703d509eaaf6fdd82be212046bfc1fed8a6f09a0fabe18eaa18ca", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "22ea0c2e45c7e95071641bb533a13410aec829b474f80362cfdbece5e175d85a", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "29d0f486cd6489f3bbaca651c0928360160c6af7bf630351464fcf6763d109ec", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "14cd255dc5b83fc3653c3302f0542973d93a4d241a00c61d0b1c0488795f99cb", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b140bf8fe02efbc690f29d81ec222dc5ef9c8fd28454537baa984dce0845e2c0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "f060c667524503dc555ee110b9fc53941f3f254395f08e3bf0c1a5ca77df7368", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d1423460758b7badb523fc8023675ff6aed748d7f8aa2e9dc2d46d975281dbbe", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "dafe0d32a1a6e1d0758b1957d570ee5f83f0abd83a4a688cbeb3982c7df56729", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b5c768d17d794b9a0d70aebfa3e08237aee2fc316bb11d1b444d08521adf2f08", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "071bb5f051a1a761617882cd8873f69dc76e8f2b9c08b439087203e14d1c2a17", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e1747ba90a8108a72040c96949f9de2a1d2e2ddb6d8a9d2c037a423bd77d67b0", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "a79446183b1f2d901df9195ed033fb4a854fc600ac5fede67d06213a46e70148", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "4fa3f998f046f28b1d6a72b75e72014ed9b0c8b4d78bfb70ee281d46261d5c3b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "450dbfa92ba2b431d9c11f1e05754c0cdb93e4965459952fdd8e25f4fb17476a", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "c0473214e30eaaf8edaaca10206e982968f23da3a89f0f2c3e94883c9ecc5126", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "ddd6e5d6fb42dbf0709e694d10cbf0b5756c8afcb1edc8814c1f6e6a81268ae4", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "3820755fc12b4635984b1cc2f999fbfa25c11a8c5a60f7cf591162daf275dc90", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "2552e48392c7f393e3a36076eec880c7cb88e5e80074d92ce2c05bd9aa49cc13", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "3c2ff60021b9f0903e544f4ae7720b527f6c908ad2889a5aafb8c265be1dda5f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "dfd6e6de2297d5db178193e3b2a785a344511c9f29747a2065f45382a33eb942", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "ad0564b7e91708f75eb9ee23ba70e2804f73ae732d140bb39108446946e7a288", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "87b3d852a8ef3a7a3af20dd0b65c6f01b01685be775388db8e2e65bf6493de85", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "676012f8a02e7ec58fd24a2ec7a331cc220a2f3f0de743153049f80220430182", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "f4ed147acd3c19ff6c07a03907091e09f1adfe48d2deeaf269b3987ca0a6c341", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8b655f44ff7bb3be715bb69a72c4a343c03ff0320d7dace4ad02e21f30ae0e51", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "5f2db2d640792d423e64ad637f50fb650509968b0cc678de4d00490f7221aee4", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "88910bad2cb5b620861724975c47f095d3cf4bf394edeb0074119574e204403a", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "9d455dab6cafe6bc39e0c3bad254a9a24781bf84ac506b869428b8f041b8966f", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "9e67a7d0d4e2508a4d4c339ac4c079d1d5560d17cd4b3b900eabc4c9becb1b9e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "cd75deb093f87abbe69c1107ec68a7e2b98b624ba44665ff20fcfdf3b9286a81", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "14889bb9049db44a65a686be5c6c01ff4f226e5243383d9f818d7b3dbe3ba99f", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "3bfae0770605956e51198d5e8da0967b088122c671e751746236aa661a4cd6d3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "063c3b05f42ac284b1a74993bf3820ed916e21995752e778650ee2b3c9c2aa20", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "e4900ee971fe0cbdf103d9e45f7cf6ce7686eb67f16806e822808c1e727cb889", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "685ba6cf7005df787298a41d35ec69da13d2cdc498ba1760e4be4e53319dda83", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "8741f41da97b3807cb8819ee2a446c5a23c5b9a5314b77d02dea9455ec73c2c0", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "0fb5fd92441ff47a325d06a902be0b03ec442d842a5d29bc8a05fa100e4c991b", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "dad7de2f4a3242164ddd65806c92ecac0fb80f85326c91cf02b0949b45b4a430", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f25dd4954050e51ea52596a19a4afd3344a391eac3510f42be2bd38f9f34c2ae", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "0bd48745af6f10a8224ff4a31de394fcb458168d86ecc9cf786ab22123333499", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "9e883c30a5c4ee597fb526cdade37b6fe045cbceb2de442690a51d13d989467a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "7b481988bdd68eb458ccdd5a61b484476a270c41d80eb1cfdf9400b4323fc64f", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "7903c04a8f82b8c92f67dcf99d8caa736a3f10cfff331ce6aecbb0fce5c35666", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2fef2672652cc5c5f2e881fcb7b313cc21cbe77fd8801967ad7369d42bd36617", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7cbd6c9d72c8a3be70544668a57529b4fa3d46d53bed3604f5e27534605e378a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "fb329e51aa4a15fcc8ef9796843b60b8669d48b17ed2aa2614fe94abf96d9f04", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ab5987e5414ee5d822204027476950a86d4b2f111a4e0005613ae003c5a03d47", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "a544cf1aece0a0b45626eb6632a7801f151e6cc4a77da39957408685e98b827a", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "e96e320bba0ea79a75e13b48d9dbb6eb1c95eabd657616e8b1636e6fd946f231", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ff207427ea11e29d96724a2e6de208c0ca3d719d99ce847f93059add47db8fdb", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "8abcd6fd49e4021585539a6296b75d34a7438990b8bb3a4db5b642c35d23d6be", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "cf6972b0632af625c417f317ab005099471805af35803e4e34a6509957e39139", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "0d508e4d1976fd8d1494c6c7053cd203fb2f3dc3901845aa9e70219b468c9729", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "8b539ddbc3fb5f357b41e7f872ed01e80c4b10f8ba8991fff6741fffdac66003", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "a40dc1985eb341921451d60e72a2b2e1a7102d90fc97bee40207e4ae49a64111", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "7265d722439dbb575d1e379b44aa3efdcd36460d79b7a3039c9e9b724ee0677f", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "d87dde418fb8b29e0df07c44ab478d1c657492e3b3b6ec804e38e234f9f3ad83", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "42918a78b5fe9ba8a8bfc3a6d642720c0d71a7ce7d574dc2d3cb46dd76792c07", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8d4afdd96672cc7742936c1dd185bb4eafcb5a86a347642f5e1182c9489da717", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "b5c85818cf9f26a4341378792adf31f05a45b0048f830065745b758528b824d7", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "24bc260bb30132c106501e6ff9d28a49bac8292fa73f951c0b615758a2dae8ed", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "d690fe8c7022e2f860e82f31dfed9627462ea79bc4baaa29443c0f5be454fa32", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "76ab5c4d45dabce8b102562ce12a62323ab74ee4bd0fe94bd7454a543a9e3d25", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a0595cd87c174aaeed7b6b12935102fbdbfe619e071655692aaa452ee09a5d65", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "c5d19545b5073606ee776555c2777dbbfd6f29869244d83b02da45325dfef285", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "9fb29557134a0511aa83ae64f1c6953d94231d7d9d10bc1426ae953349556d0b", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "db0571a4d96ede1cd9bf71e0c6abb29901a42bcf0fa268f8b44c372594dd5e1b", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "1752a61c218e78bdd848f62cccc8d330739efcc83693ff318dad31cdb55dc88c", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "180673a1014470714d0824a6736bc75516aa71238803eea3d355b10fda67bf66", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "675d0fa47ba9b9cda025d15ba2baf0d5058b77de2777084ffade1b2f8b210b43", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "962a1d63a729021e1d49fd967df5e1d302d5087efb89b6ed1d7119a38cbf328e", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "895bf325914c4f5cb6c0913e07cdb8a5bab7259008e96267f35181214c33ae7e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "fbb723fa449270f9a5f7fd677cca0973eb3a5bc6fe5da050e90dc32e626a4b02", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "8145da82731394d4efacc359f39e0a6eed62ec2c646fe5b2cbc6fcf839ab8e80", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "ac2bd4a713239a7f9e6c594c317b39f8968b276dd98321bc1ba2b12a10d34bfc", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "b4d9d4f21c72af589b791687f542cdd268aec6fb43f2273c66f14813c3394c75", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "1887286dda495a3e269e9e1c332f95db3de35664a572a05edd28e1485a20d9f7", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "b02a9ac754fe5e84c316eed4d953d13a6e84a567318832ecd5b88f8d1c5cc1c4", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "ee41b2f62421deed7c537c6c233c750382e4ddd47ce1afbb7034b613f02a6984", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "35bc724008fc7f5c3cc6e5c949b6de09b261c9822e9e93274eec561abd783be5", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d7e5c7948cb8fd891bf38ae53b0d54120190026fd6a5a5344f9cd135d6352d88", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "52f1c445cbcdee55eb6e3e43a52f249e9c6b73dab088afeafbc901b1099322e7", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "6431aadc0120ae92f299fb931f741e6edbc057e9c21e31aef3d01c5758f70e0f", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "abb24be0e2b60220446bbabdfcbde0da60bade05ab68bc8e9f6963ed5ed5b6ef", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "57952075519e18880f95d7c918dfdd234d88a4eb8ef0176d4aa147fd015ec9a8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "11ae7038ed922e9530568bd33ec6b4dc45f5b1c06b9a854ac4596ff73050e5f1", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "abc70df2dc768cfbe66014232f9d4dde32c66e67131d71ae9ac7380b7eb17a9a", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "9a2857faaba162c4671e2410667d8e4b7930f55d06d9fd5677b69ba6609151d0", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "c02d57d8dd263a71eb4b9d7a64d6259f4667c245d3a723b5f441cee85113d1b7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ddff459634b3c997e1a42a40f4b7d8aaea78ae4b3232077276b1ea8d7934cac9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "2afe14c124daae6ad48528b2712b9b221ea7e9b17fe9b7fa1acd7979fd8439f1", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "bedc8094d2c967c372ea34d410157389450dd6368c693b0d40d2f0a116e52cf7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "1f61ab37d41a19d7619efa012a8c4b9cc2313b220f9e9730029dc2f065576ca2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "570529746f6fefc2247101f04dce48f8980cb0940eef16b3f5752bf07f1eaa84", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "7799aa237af1dd495ff36a5c877fa04448cc9b6bd6f9aa36520fc6af363869f1", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "f1a3214e694d53bc53373c57287b06769afc9d419c91662f3cdc71f96465eb12", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "74f2ffdb9c8576fc34be860cd25da1a92dc4c95324cd595cd92ab66a7df8814a", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "5e992eba0f2f41c6cc1690ec4a4b4e74f87960882cf97a43372217388caacd03", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "ba7bc6fa3a0dc6b5c3a59d6127df6db57a1c6728937377db8fa597fe8bef581b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "b582eeb24e2434c7bc6eb1fc88f2824e251fbfb62c76ca427fb720f447fc52c6", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c2220ad54273313c8e835c5ee9c200e8e1342c4e033f2aa4d9ceb4155abb241b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "292961ab5fb56e414f97b989fecb1aa100fe115f2835301c4d04760c698859e5", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "b72270e4497526b515c2129f14be20d5e79cfce724f8b4042a1d18a7bedb1d87", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "7a1282d59d9d95b08db19a0e093bea4b46305733a6d99a1ea78a5c46e30abddb", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8586c74f493739486d60e5fa6ca9ad14f82d1ef30a36d64210d272cc3c194880", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "9233d25df6862a7a495a480b84f26247032483b24a66c763506d2a2f08d80c02", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "673ac0a4bc720404d42991a5236bf1d758370760ac4d11e024c326e6f109b1a3", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "7b1a558789cb9cdae35910bb50365f96d158cbbd0965a0e6dd87bdc1c9b5a131", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b9d0aeeca3e5e1d01c1ebdd8f882343d216a799f0e9065d8decc5217afb86ec7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ee31c3014cc1a3fdf403ca79217d87128f82b69ad238faadcb00e26a439ad0c6", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e8137279ec8476615d29a333d4faa27654cf87b08e91a7c019c775f975a89feb", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "5aa7c69ceb6b0bec3096092a467337c6712c31fa8906a2e6decb8d537f2d3c46", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "60537b479a369de2c4b387764a5507cadfbf6d0e70df0e297c4a3df1866aeae5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "07c00abb1fd7f937cd7482f749d525f5d1b008a628b4d3ef4428e6b044e2317a", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "41be9fe9631c9e2c633218b38fcc4c447afbe21f0b2e5c5c6f78fc227ef78eb2", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "df1d17b942751a80afd96872f5feb6cec0aa61b29670759220f793baef8f5f83", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "cb12b21534672658c7a1c467d1782eb786f72254a96034d66110d77fc1679924", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "eb59fb3a055fc464263386c4e2efe59c67c9e8488bb30bc1f0022a1e3ebe709a", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b31cdcbbf201dcec48b0e507ee3d277dbb6a694b4e5a3cc7b04dd75183c1d38c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "fc43d2cae7690bc8fdc7bebe8cea65cb358295b12496de23a4d72a0b3678b988", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "a2789eda1495c3aeac93243e9096661c781c0187fa2b6c231ec93345bf21d89d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "2c60f3c81f24211d01ba9ae182ef1084dd62bc62b4ab65f1d2b15b5b0c0fe644", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "fe9e4e480c70eb472fabb61d4c5e16dd87483f7c030925b8ba34c6d6dbe8e597", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "7ec9bd5d3c6522c688854988d1df1c912d164b8c5c60ba90502e2ea20f13e498", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7076010ee5b93a32e1717484cf9f7e51bf679869dca16699d542d9e5ec5e2225", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "764510a1b7239df68d3156b985dc6c4041c1432b58fc46ee0e2683576f683b34", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "7e4d3671f9c32515b6f9686f51bbe5baaaef426d754e7c3888c60ff93590ee93", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "bc59134054a0f68b405a821d21cf4adee64a1187a3dc00b6eaa9baaf561354ec", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "638096c5cc52166ceb33ef9679f249ee5f8087415dd6e5a8d8c6e2a1d36f2464", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "b340596cb7cb45b92eb01ee14e22887ae04c7374a4e8cc5f9d71cd6fb415a50f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "bcd6cc2a5bc7d076467094ebc59d47f1f7cea1d1221e4bf73bf55c4f07ab71e7", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "6c67459199aa515cbbdf8b40614dfab604aa8df11f7330b6ebf7ab4a3170bbcb", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "f89d599c894a3bbf2cb5e145477d71e31c4ec82e452741166ee8452d1917437c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "70668530e85bf6e36ee17850a6cf4b3f8b562f0e7879bcdc80876be017324cb1", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "33184b6336fba46912989b0c336a1f35637b0b7f416a8958429e2d539e8e7a4f", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "30a0ca2d2521ecc242546d4595b1e431fc91b8e173c83fd39c1887c5f5ad75c0", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "70e9fe7081322b944c0a3b22b12583d798b09ebd1ff030b634552eff8e18982a", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "a940f86497e22690be32da38a42268c93298d56c6ef67f5261d7cdbc9da0071c", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "67ea2b39dcc46baf84f69f02b1082282bbe5e3b78423335da5c58e021d7f95e0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f04db812b75666b2e4409ac09a789d2c443a068d4d1c9b17ca74112fb0301952", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "78850f6bc410f4f61d703833805f27063c3423c15d5eea3071e88b2fef55c481", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "b37ea46297f7f3aa6ec643363402352b7a606566df609acb10ee8cac934fecfd", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "153647814256485dd540b0c1c5f03dad35cfd08b2b03f362faca2036a4c7aaf7", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "41f5b889b332249d24b6da3e54329630a2b1ad33b4fc9f492f0a083bbd3e7c98", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "367ae9515caf718cb377e7b99e841cd060c7c5d62d497f02fb4f61ffcdc4d899", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "374766d065e39d34db2bc8d9ec0aa9225000a83d7008102dd55814021e0ab9bc", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "25101b908e805ff84c5ff47eff93e75fced2d1ae0ae430e1ff1f4639a697ad47", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "3880dd70fd15d2aad1443441b47805b0b8cc3bc6d9bd5299f914db1e9f202dc8", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "892e0bba0697d79e806d8309482f2d67b14c967733d4cde1060a43e22d237f02", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "6314cb11efe8702307f45b3c23d70458dca4061ee498c798c8ca5b2f5b47844e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "733666c0e5d3a948fd1a998be53b2e9ab17174d860176e0e03ce89a0813dce10", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "f06f30cfcc60e73f0a9f6eadb83eb839cbde55a549f1447306cd414b575e96de", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "4773b6ee52744efbb8009d2e49e9e954ec1f300f4ff3c1d79cb436bb48265aa7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "8393c7184786949f0f44c68bd0733df61cea34cf5cb35b7657e7de6aab4a5003", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "6b7dcdff6d495d275d83b11e8dcb77fb005ad908bc05bf56a4701c171f81bd0c", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "adba66ff43012c6baae5859f2b13604e2780adda9c066598815a0243c45e79f7", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "a95f1cbb0db4474c31b3f8a1af665814f9f913ab072bec1ad80b635cceefec3d", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "385ceb6a33fdd14af87fa0b14084d817d070bcade75461a8c8ae7284dba5684b", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "10a1377e3c85cbe943bcfa520ae589687e469024b52e85f091455b89ec4d9e6e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "58bd510a4237a6a627489d0b90b649c029e7b19a4b70babd9279e0b75e70f93f", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "ca254246e7a1c8659827f4d34e3bd6850fbaa02c3c088684c132573d176dac8f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "a51cb2f2e26b45eb71e7b1e170c9f4bf8a0ab25bbda8efd900b4943d0f0cf946", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "75bc39a253374eaea7ede2d44f70b1b1eab376ba24e15cfdc3a106b717f40cb7", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2fccdd68e494f3067843e14f8850f4fbfdbb908175697b19a7ad16e90c95881e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "ecb2c2dff568891215bbd632ccc114c5f50a738d001b4ab0caf8348bef1ad2c9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "b85fdef5c84d790d08f60246e526a593b11e7899540525ea4fc01605574f67fd", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "549cfb5b213d9ea7fdfc36254ba68395391a5f851e98105d12e3e6aa85eac4fe", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "b217a365ef35faa53158f72da165800922db45ed7e8ab63d02942efa20307d13", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7381fb8b843ca1192f7f87c516767167ab10735d71038cc1d37b73e2670c5ec5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d310678389e1807b9cb98cf1aa82c4cc67f2cf90c78a5db7217c9641a9ac4c54", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "acf300b6b1d8c07e551944e233b524a3ae5b5b312055921a27bef58d1210fbc2", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "8e07daf3bedc106b836266472a62104b4dba93db4c28d4c19df48c3e160b885a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "8ec604c60fe0129a2eb2e0795282a59e108ea83910250f6396640d21eabf3490", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "1dc555ebd01f0d650903678e0f2b4b6e20294f1fb5f4f811046f5b9f09978bf8", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "6198f3d9bb538d86fe6c9ea57c3ed9d1c63d48d94956b7054a147c8951b3d56f", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "ac43f3d213bac5fbfa32d22c50859493578a152090283035096a858faf523444", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "d2a2e4552c9e4b3e8065bd175a512cbffb7e77d2b726a53ee5bae549d36852cf", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f24f57aacb5ff39fb821a80e7674fd9f60a16e2db7ecf14aa6a94c909ba95d9c", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "66d1cfe723d064315c92c6fdc9ecd7e6e63183c67a8ca9b7e92e9faa14d0a6b4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "8afcec784f188447218c055f74fefee2fb45a620e9f0efe1d6f2ac1dd7887834", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "7914a32358973f327b7a374959b39a6d76a0906063abd06dfac7358fa240538c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "12437c97636b7ac5968a7e366325ff0e89ea8735e6bee9b81f2aec1314df93c5", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "09b54f137ff93b811570d29a4ed09cbbe2e658ad4c3ca949c7161c422465abc4", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "13a0422617c9d607219fc496ede524d3bfc8b573884711fa8a9844825e2ecc17", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "713dbf9deaf0c23cc6709b101606164bc784e00e63af1d504458afce2654e258", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "7cee216d71d088ae648cfd9233935492de3041ef9430df0ff2b28477c7ab9bce", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "4ede89d094a48ca2f4718a9095a9a007dff35aa14d4f4634ccf3af8f40004ee4", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "691fb9b54bfa0820fa36f05d7097425204fc883f41ebbdc6371f8649d4a92b60", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "38cf5880e1b544375ad056a212cacad75b2cd8f67691644989bcfd184b198e4e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "3c4bb26e0830bce137717450539f1e501fa817a24625b416bc2998cf54d3da90", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "caef4094191d2ff3e6158fcc21ff9ba67e5d4e693fff97490f87039f960a12c2", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ad16690666b9a76d8ef86bb20d0a679376be315a640993a54d322b6aeff8489c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "eccf177af806242bc9035cffe917c0cb68c88e944bac95272c697983596735c2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "d13bed691e72b3cf158c55c9d4c17c25f4e15e8776420b2f970f061912bfc164", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "2bc832825327ca23430a4e874a143da743167495d1bdc74256e8f4312c4a6f7b", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "a8f4a2645c81004ce07dcce9ce0966aa7e19093a075241656ef4e4878c038283", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "7a8c41244d4e1959feb77b56e424b097349c0a4dfe6efa627617c6e4cf623903", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "015fc313ec8b2d946dff5a9090475d5c82137206a73ea6ed562c66befb5d33e3", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b826ca06ba17f3b9256c97f54c91fe0e5f00ddda276825d457319c6d046a3fe7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "c494ff22fc527f188fcb2bce52fe2d24fc9610fb8af8dc66b85390554af3779c", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "a0019aaa7f6632bec4aa862f00b7a7643b06fdb9125c87c87d2fdb32758273ac", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "d04374435640e34d7b11339710924b11a656364f09160fe396cf7d181bf9784e", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "0e191a66546c73df55df84854900da368f24cdf7ad002f96e8ff5a428a48dab9", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "efacc1b80daddb2b3721d16d7a7587a264d9d4d63a92ffdd42cc861d21daaa70", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "ba7936202f01fc9790250d85239b6cafe1288d91434379079cc6f045db89d591", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "321b6da2b79f5c8be6a0eac07d7b8c99fac029c86077ff9d7db021c2017255ba", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "2174b9843f7661fac7e450ebae24b69d5ab244cfdd5d5a30c510385f418aaf3b", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "cf6680b17ffdc68368b5e13607becfa10644098dbf3b41bb7713dfff06ef324b", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "e5dd28594c8c1f8452dc174a011f5b6b4fd417b2f6d6e34caed867187c89bf36", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "65798ab14e6a78bea4396e5c76b6cfb74776f96a2bd564dcfcf09164ecf9b494", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "c6e555f5a07fe7b98c668017987a3be72b7d91fef7d9a7a15437e7d3db8cdc7c", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "bd6ee6b87174ce8e71ec238546fc33492143611891e82712eace1a24c965f62f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "8e2e741766cbfa54ebb57310a8c493fea2cb41127596a17bd77c665efd413334", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "a8bef23f8aef8da3e7569cef2280d16983d1a1616b30dec4e7c531ca733bc733", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "3df052f68f4f41ef1e3cd87a5d7f209ccb220d668b9db280573afd1698abb109", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "e036e5264cb785bd918aa63256fa8085146b29571fdd6a618a3a1ad8a87335a4", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "a0a2768b10ee44b34b7151216755cbe38cd34684a7d6573ccc5100f9973d7509", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "11db402c4dc7a580f35e60347ef9cdc7f525174de289793d6c21ce4c30542d99", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "4f5d9572a89a82a7e1abd089255da1800b57b9228f4d9fc146f5e90d0b7e9f19", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "c28e277d49d19a6622abdabeeeac68909b38ca4ab20ff293ac6fc17584504bcd", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "99b8437f1e43412bd281114602ec151e365aa9969529d14c990e3ea5a2fcbd88", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "ba2915987dc00ffc9785a5dedd0ff007f769d8ff1da8721e09965fc02c7defa3", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "9284e933a2149e2ae88f71784fc2ceeb7abc00de5fb9a74fa3364c23baef78f2", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2fcdb2256172efe63f64358a92cc3f517dbd68fcaac9b2004356804103413be6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "11f1c32850edfad3490b8b757c3498f064cdfef70024db8df54b3c10d44192d1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "467fb8032e41556906e06af65999e556930835dbaf6bcac5db54f87d083fafd2", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "a9ee5269d782cce8d56f95d1d09acaace33d5913f5b4d65493b053779e0955d9", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "9c602ce85f9e3a4e7f093864f6f9afabcbdba9b8b0c14a2714805a3436354545", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7af9a5d50a0121f0399e5b0f52e3a35000c9104da0aef912fad23ce47d727be7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "6ca9f5e18675cef719c4631b7652f275635dcc50d03e1d4ed443e26904e17c84", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e522298dc7ec3bba8902bc1607b17293e40c186f5e9ec16e77aed88ca7721824", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "5a4a0156453d2e950c10ec9b69ce8d05a12535ae260030c49f951d999a73fb42", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "19e55d732a19903f320e19bbaf3337a96bbfcbef3808d18b591732f051dc8fa2", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "16f2d8f7467b812157fc3f6473a12744d95e06fafd10d5117bfd8fc9964fc187", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "aa1741e5d8ae545e72a4ef185e2c8cb53ce7f0efd1b542e5ff76a052e5a5b33c" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "b5bf5462efa3db7c6d6bff029a9df2fae6755f1622aa1a39fb2e8a202451b411", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 26000, + "training_contract_sha256": "e55f0cd01870441f33c9f35544428cbaf029edc34fa1637973d29f64a5242966" + } +} diff --git a/models/ppo1-random-walk500/0/_CHECKPOINT_METADATA b/models/ppo1-random-walk500/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..4cc716333949e2a8fde91854ecb39b00b560f6cd --- /dev/null +++ b/models/ppo1-random-walk500/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786225531877699381, "commit_timestamp_nsecs": 1786225532597542199, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-random-walk500/0/default/_METADATA b/models/ppo1-random-walk500/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-random-walk500/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-random-walk500/0/default/_sharding b/models/ppo1-random-walk500/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-random-walk500/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-random-walk500/0/default/array_metadatas/process_0 b/models/ppo1-random-walk500/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-random-walk500/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-random-walk500/0/default/d/036e4c234628f458d5f7a9847cace240 b/models/ppo1-random-walk500/0/default/d/036e4c234628f458d5f7a9847cace240 new file mode 100644 index 0000000000000000000000000000000000000000..06ab320bd0c2c2da4950b77855b6e78ff79c4379 --- /dev/null +++ b/models/ppo1-random-walk500/0/default/d/036e4c234628f458d5f7a9847cace240 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:04a1982ee5c5175ba46d0b07133813906e48565d2a91341ba1d47a2220e1d32c +size 112833 diff --git a/models/ppo1-random-walk500/0/default/manifest.ocdbt b/models/ppo1-random-walk500/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..012213eed2c127491d49f3eca3b273d472e0a376 Binary files /dev/null and b/models/ppo1-random-walk500/0/default/manifest.ocdbt differ diff --git a/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/13cdb3e1f6a7eb1b28d75f2a4818f613 b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/13cdb3e1f6a7eb1b28d75f2a4818f613 new file mode 100644 index 0000000000000000000000000000000000000000..e6af2962f5caba1e866287b1d7cc9482791e9bef Binary files /dev/null and b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/13cdb3e1f6a7eb1b28d75f2a4818f613 differ diff --git a/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/2077fa6cda20d4df8af93326760d3a2b b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/2077fa6cda20d4df8af93326760d3a2b new file mode 100644 index 0000000000000000000000000000000000000000..8f7d74cf5a6afcf886b837e81e86a7bcadcda370 --- /dev/null +++ b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/2077fa6cda20d4df8af93326760d3a2b @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4ba688a02e75852494fe5e22c716aaf94df055bce966c9f8dbf33c9d747e3d34 +size 164593664 diff --git a/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/3cc37a3501a39d61c841c613b1e8e8d9 b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/3cc37a3501a39d61c841c613b1e8e8d9 new file mode 100644 index 0000000000000000000000000000000000000000..970e98d09ac07bb679ed747d8bea6c27ac253f9f Binary files /dev/null and b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/3cc37a3501a39d61c841c613b1e8e8d9 differ diff --git a/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/af9775c73bcd886b3f09f7387af537eb b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/af9775c73bcd886b3f09f7387af537eb new file mode 100644 index 0000000000000000000000000000000000000000..4b1a203a494f8d5cd524894f631d64fffb91355b Binary files /dev/null and b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/af9775c73bcd886b3f09f7387af537eb differ diff --git a/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/d9ed0ffbb5d2cf78cf8a567878327ad6 b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/d9ed0ffbb5d2cf78cf8a567878327ad6 new file mode 100644 index 0000000000000000000000000000000000000000..759ed8d24465838bd2a9c7cb60304b578288faa8 --- /dev/null +++ b/models/ppo1-random-walk500/0/default/ocdbt.process_0/d/d9ed0ffbb5d2cf78cf8a567878327ad6 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ad6a9ae88893b28e6e0317d98407123fc05c7ef6a9e71ee7f491420ff8b29089 +size 414674944 diff --git a/models/ppo1-random-walk500/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-random-walk500/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..04921489ea5b39ea75af0bfa1687ad03535455d1 Binary files /dev/null and b/models/ppo1-random-walk500/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-random-walk500/checkpoint_manifest.json b/models/ppo1-random-walk500/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..98af96636469600d08dd9050510fc3138e113cd0 --- /dev/null +++ b/models/ppo1-random-walk500/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 28000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "b03cb1955f9a851b72bf18139a61c06e791410f56e3f417061e706dd2d15d1ae", + "runtime_config_sha256": "b03cb1955f9a851b72bf18139a61c06e791410f56e3f417061e706dd2d15d1ae", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "b03cb1955f9a851b72bf18139a61c06e791410f56e3f417061e706dd2d15d1ae", + "schema_version": 1, + "training_contract_sha256": "e9ea284b76e5ad80bc54e88bebc91fd8112b8fa0a1b1dab774c4e2dd6c2b8718" + }, + "sha256": "0878a5415b778dc826f3f0bc2f64ccb67b6cf3a55de3c837822a27d276636ee0" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "073981c76d49700c99ffff08a6e45e812dae7456da311a93e4663e85c1204515", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "e60286c06449969f941ec3d59dea06cfe0e0759a0977e0b8f2ef64b6fcc50157", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "0dd2f2f3105a452194ca5c4e9989bae6244566b9bfffd0ac833c015db7222e19", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "4381f0003f1102e359dfa72ed62f53b069a409a5d8254f9411ca429c2f766283", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "7db9449e86c1c874a61ff757b712a690ab6a9c92a156795a165223c5c39f08b3", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "53383befe05e6ce626b9af008a4e786bc3d7b2ebc52f5b4541aad92b8375b180", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "877da0ca2069285403f1278ed2cd5b702e3f6c3acb679aeff0639bbe0ea04bf9", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "8cd25211c4c3ca88156ff2c1d0b5c0e744dd868a4844ada67f5bab6d9f0c5dbb", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "10d85a9ba560122abb39c4307b9dc28ddee990063ff85f322bd75400487b476d", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "713f9ef16c3de6199374f1330b6033c462fabf719fd695ee16a1bf7bb95633ad", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "d3bf4cc4dcdde3df6783e47f2ed32eb7fd82e7b6761533fe27e274e43cc011bd", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "e4f83ed83075d6e4430adc37786fdf487b293ea74c9d452ace66ed6da92d0a86", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "d8199c38a382483d339a42cf9923ccc7df53d3d9ecb8b00834e70d19f00aa668", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "4fd0dcecd87be6dfea41a53e5c2cea5ba8564a76a5a088861a5cf4098aba43df", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "984ede63b47ce4bf812397c3c2af9097fb63c2ed221bcbbcc60cb377e110701f", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "85dd889c027b52e6d95f370ee07a71f4c20d7ce2a4fabaee6c30382bc3bcc959", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "c834ed2d9c3efd3fc41eb59188ed9fa2e9b2842b207fda719193dcc3430b6a99", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "dd812081360670f5cf9ddfbfa0612d821456e3ff00cd5594645022acead029bd", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "01e57a1c2296f2874c64686a1e02651af40ddc53963cd241be1f618c3f040b32", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "45acc3000e5c9ac5e0f5225ece4019fea0876c40e3117fa4b0f7c1eb2dea2c97", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ba685a47af8ce758686c9a730450c74217d86b9347254d91b66fdb4187b24dcf", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "e48c48269a9ba35b2a5c11a86f8f8cf0b957bbb273dffc7cb6453b141945874a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "3fa5f607685d7dfe32a3eb4bc3b274af53e5ec5ec60bd9835903280baeb968e0", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "3ccd22bb6b6494055da04102b331a134c644ced70e26d1742f41dfc2864e1dc8", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "724d9d7926d34f12c72e33739593ec16d6b5b47c828553cff3838299d85869c8", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "bbd3af0eb425303f8c0b197617535e20e19077d050683c5350b58f374c9b0ba1", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "20e208ee641c4cec20a71b412b5bc8c2bfe507dd5685622f95ff862d0543f372", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e74b740424730a363b1161765037e02d0659a3f54bc63d28c617cc86e27ffc77", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c344a76183c9ec4352c22e3ddcc42a244d21467675d71534f34fed083704751a", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "cc0e6fd425aabcd4c139bf2924711e1d6442e2416c001733f6deed0c20b11d96", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "a5fa1e16fa48d70e0dc0b01f384c62079ed9dd83b1eb55048f3d78090549ecc9", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "456567b39dc462f836737e02312918dcedf0844e3c7f47b933a67dae324a9316", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "8fdee5dd48539efa74fa6d1f65f096501f88af2787d920e46d771d3a2bc57484", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "bd64f9e9b2dfafef9b5bfe283c078faff7eabb2c8fcdfa867c659ffabaa157b7", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "8351500fd606c0fd9955202ccb785283b201a4d4834f6a0a03dc4024dd673a8f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "696bc8afa319c54a97fc8a8ff42edf254c9ced8458c19864e81a07a4bdd37519", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "1f053596ab11c61a3c809cbca66d720a25d76cd925a4eeabccbfcbc22e0a00dc", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "5dd6a3ecffb02bbeda1457b3e5b45ce83083dfbc34aad4a76a6248bf5f1f18be", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "91591491f16813ae2acda2b2ba9ba1c267832524fa8275e3e8606363d29ce072", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "995d1d7bd972ffb68bf9803aa3623ce72070c1e6ef0b86be058464bd951bb25e", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ffb51d69b408969667ce6a329db550f5616eaed55d529e4308ed35e29c0bcef1", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "03ec7191c91d2852825cec026273f93fd5e8abc27aab48399490f35bed1eb86b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "ffe5a55e59e2fe077059cea0245d311aea2c7c41b5081152485d29bf6736aff2", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "690e91aca53c0e6e5e83a81e6727eb33fdec0172135ae0f7243badb8a4be882f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "7d63b663acc9eadfb76436373d922d3e5618231e4f09d7b79bb547d95919a31f", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "2b77f5d267ed6ad25339da2b18739050b7c6fc27623c1c23f3d8a9ddc357c899", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "124bbdcb5098c0203be7b9444e5aa1b04b903b2797ed18ec3b2709d7641d5a7b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "c739c9de002d748219c8687394aaca4042b0b1fe65d06732dfeacb57bde38e6b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "219c24bdf7dad27f2e2aa63cf5443ec452ba71afa1894ee7ff69b530b7f76d80", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "96656cbd7ebdf627a8df3cf71d474e68b210a8f025aa903dbd35174514ae92c2", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "777a4ed1f0f3e836d3452386aedfc1501d6c028dfb6b20cfdb6ee6b194280b0b", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b3717dbf515609ea0c570f22104e783af58882de375d185d46139c4c1f26b0c7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e0f115f05f24faf449cb82017c1a6aa6e9cfd24faa5a1414cbbd2fbb8d782f28", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "53d79c59390cefe2729eef0587057f6889ebd5efbc3005fd3226ee6325d5105c", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "78503b08b7b94f07a3e011d74f9f27269a9d133f59dfc76118fc3b560ea90a13", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "d34caa148bac0d4ca9dc152239284d68c58a9c3351f870b594eb56dd1c784eba", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "c3acf207fb4667ac958e32cad2b52ab9e6e0d27a75a7bb32bdf271f9ee57ce6b", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "70685b3ea1e3af0c5a97db11de3c4eba2d6b78dec0b38d6766f67aa50754f009", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "a5bf7ba6aae9371f0a63d7f15efdff5b00300175e0de1f34be4ab8d48c777bec", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "a7bb0a765fb5b4b72e4ba2a5510760e95485bc0914370a67a4d3f8c5437d146c", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "5e4d6405998f3363afdef5dad3690c518a53b4372da5f9227e368249fd6630d1", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "0d39b88ccc5c92c0fb6abb08862a033f5f1eb7319f1e5532a3b19cb6a5fdcd30", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "ae6eeb6e8f136fb4e16618752e158dc1de98ba420cb93ce222504fc61c7dbaef", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "a8fad2f0ac50dc07aeee75f428e6d7333a3ff03b5b706a7604b296617c8b2481", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "988cb95927e5257a101935d942cdab861e70ba6eba3408589008ebbb3a920fc4", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "55b16949d6f84c49cf7410b1bc1b7d1ce0c0d567fa456987257632c9d84b1ef2", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "4518d18e72af74135dabb1416d03009a169fccb8599de16b2e0621035a1580c7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8d93c4ea6c2a322889ad650afb712c1f472324e86fb1cc2a3f85849c4f3879c4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d68984f95c2e4090b892ffb32b26592532b5a844756fbcc393ddd27368e8669e", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ba590eb6d80d3c5723abb1d725b212c695b3fa1746f88f8ff6f75c0c607b0292", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d104c0e4279218f52e0e68cdd3b63a425fdff4b60235647f9900c49c0bb4c425", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "4f7820fd7dded4dda8124723707f7755af11bbdb1efc7c2cfd0461f47d31f264", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ed45f70b40acb3113891c8f2bd42d3296e8ac935e26e8f93b66df9ab27c4a152", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "e92b4cf310bd497188e1b106da372e176673fb33a3ca8dbde83093c5457310cf", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "19a3e61bf7f73401b6f0bf6b63c20c44fc1d84cb2e5e56965d630d3f42e37a94", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "b86851d3be7b597fc64a0314f2bb86826a7a27f335105e7d0ffa9d8504679a87", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "862016425aa6b73c0d032f6364752c6cc39234682aa67d27d98c30bdc7ad3ca3", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "c3eeb1d9c37ba3bfe344ce6a04c89d09ff7b74ecec5676c11ba4eb09699c32f6", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "234d3fa1b90655a637a4d8dbb985c6e48f8520f0751500052626bdecdf111e33", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "5098d534742e57af98d675c400d4e845fec5ac07244a94953c93c7005a7c2636", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "87f9d698a013234f027f60a5a6aaa7842966f02161f3dd627ef8b3a212aa90d2", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "11af48db54e2591faf91406c47b158ee83bd15f109db4eb6cbfe637411264363", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "5d220f7464b590d48e1bdfd248761a1d7f35cec4911fcbdd310aabf1ad96e5a5", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "531d7cd1d476bb68a19ee0b276555f8a94eb9fda2dc5aaa4732f6a5840a1cb94", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "b48f8c2ab374dbbd9731e5ada5e66b102d4f0ddcbd84c150cbf5ba7858608390", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "93e0934bb62c65e0446b7e59a2ff58aa8ae1f600818139cacbc6900c6891b037", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2b91951bac07dab1ae5c07ca49372e0ba2cbf45741afb6f03805cc1727afccf5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "9543b88fd20e41115ddf70f905eace59e0ab5e591c5cbb01797556ca2ca1b25c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ef60adf0be39315bdfddb350a5c2f1d3c316aa9a3010ceb3998c57a128a1b572", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "46c8bc1aa5a4b3b50543bfa5deea90a45a85ca710ead71c220e78fc3a5f8a21f", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "6e2f6d6ae5f9346309e6f3fae88b76d2b12654395be47e4b74cfba060e06d817", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "d9669e6c613fd611d59f704498cef89c73358b28ae8a8cf3f61d790edefbdbfa", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "fe6a31e5badf47ea6a80339bb99c8ab673f62c54778086c8b35e819820a9fa27", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "dc7dd632a2a58e538d3839795973b4c156c07372f1e5eb5a5867e5292c6338ad", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "4d2e7bc40776e26771062830586bab3a0b49bc2a7c650f235f5dbf8e3664143a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "31c41e709515330b2b14b1b82053d208094c90d9b39d5b1131a4a31a38aa7655", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "d367768848806a00a4e50bf9f654ceb630a7b459ca750c3cd751ebadfa5ea70b", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "396f86d2e50f016b94590e3fe6ebb28c251256269dc111df2b3d250ae596aff4", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "f7ea6892f40a0644d7b1514566dd3f4151bcfd925ad7f4030288eb46cdcd49ab", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "323e022f70fc0448eaa2bd042d93f5b80b116713ec36e67cd9716e47f7d4059f", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "4b0a238e1f73c8b1ad45fa6d8255ebb4446e8952f4f1b91f597039bec1f36526", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "c90da209b740dd556cf2c055bd936f12e78fe5376fc35e9a6bc2a57eed4e6c21", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "0299215ca372077e4638fd0652cfb0115e181740ee253bc2a63362abe7ede40f", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "bff99810388c7b7a18895f7be57d62d4e02617108da35318e1af8de4eeea95b6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "338b4433f4077ca2338472ca30591bf5a376b43efb3e0edc28132523dc739cd2", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "7c96fa8d17c7fa4bbcfd1c81bf3bea0929f88987e8a8a38ec6935d93fb18ee61", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "231e265f16f196ba3633da367db01b40ca06826bd675634a7018d80e780709e0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7cafefa9ad195d2d2f31119373a92fe37f144a8a4abcbf9ec51b2b838fe75daa", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "b6468ab6d86d89ab394ec77d440e3e7bd833aa763cbd1da9631bf470c0772d9a", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "a4b193c76d72da4ea3326d989bba0bcdc1962f380b75738cf50275b21cdee93f", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "4b8d8b979ed4ee880b58808748fda155b84b119244f2f5fc90b69129bd8b26d2", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "2651c6753f79ec52ead343f3fdabf12fce71e20cdbe02c4023e8189a0adb3da3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d918334bd90effaf274460890fa034da1ec3267515070d05ffc44c6fa6f51ec7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "0a5a8a018db07ba9503749ab2caab3b7f02a555c7f2b8b54ab6639896f9155a0", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "4ada06777b8cc7ab9e9a156ada751c069dd971ad4c6529dc64dbab8595edadec", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "c40694cacf8abf60593750539c07a6a8197c22a333bc9ba5a2e31815573afa58", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "3d726042d405f7694ffe19dcdbd913fb532d92f2feb924f9c183428baf2943a6", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "8b07d7fd53e4118c4bbff7e3f651c2faac52fc34cd160fa7c20f54329890dbb6", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "ef59f7c69f0fdeaff0d6eac47b8ffd8f0b7ad9b9bc841037f290e530cbde159d", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "6d6d11a38676a077475f969ada9c377df9f3f9fb248381707e751d279f49dae5", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "72d8c13a4a3823045f69e7e1356389343d00c44ca53b12417cce779c850528ce", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "9087656589892c7ccbd57219785c2fd1cbd45101490c121fe0377ebc68ab8996", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "a8043dddc601ce1cdff8d85b1e1d28cda4cad3aa594f51656ee24a303315f6e7", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d3126e969c729a507f6b219e608795b2838c4e39f369e658f821e33cd56c70a3", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "07b55850a3d46e6cae0d7ad8499516a00f976a0d8804f6c972c4780ba63f01a4", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "d8e6699a799c35faf4693fd9b39d699648d6f0850dd64f20b83e143ec679cc54", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "e97c51703ea9291f216c1e5431573dc919dd00eeaa6f0710b67df00e68787efd", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1b1c56ce37f61580e12142b54e530279665e6854909815259d5b7cc64c5d79c6", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "3bda7f2101077bb27ce84f08ebca7332cc8d84a65e131c682a7f8089e404aa95", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "381f088375aaca6836a0ed2aaf0c9ee2522b30873bf3ea8427f77eeadb3a469b", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "2ccfcfe34a8684334322a08a1e06b189288cac8ccb6486709800dd49cc595a60", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "bfbfe0243e76f93416afaf6d52cc01fc0811375f9c8c7f6ca488d6c2770979df", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "a5dab0a4586be09fbc7f437f1f67295bd2eadf6e0eee984d8a9b2ae326e2ae60", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "39f9fdcc49c1d0352d020514e3288e98d75698087070e0cb7931b95e3d15c675", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "fa214ef7b582949da7541a29f6263dc3376a88691762d291a85546e498980759", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "fdad04d454aee5f3c43235403a41d70f0158aadfda9396f3fd38f80a835834e7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "1403aeff510ac185802dd3ea85e7de5a7bb9f617c844c888334201f636b9ff9f", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "3c0e0d777231846af71944d4d2acd287c91a5dbd665a07936e40447d96b434e8", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "2c302215af1a0751d1acac3e077de3fe46854ede8655760df387701ffa11c7ec", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "76023d3dde45f50c461779add88f9263dc1105dd363de8c403c015740a3fd6ab", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "6808b750516eb2022fb54e9757944b7aa6090e467302068aaae9568489ec5eef", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d5d0bfa9b41bd9264289d2aec3ed95b48d2b27c5617dcab53b66096d80823255", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "1f133cb64c9e3e00f844104d18058fbc3b8e007a96bbc4b63d80f77d85b4e0c5", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "9a1636b394734d26df5fc8221e611c2a36ce79f820ea8d58516a4486b389c7dc", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "79e0188c8848072438fdfe9b57e13edbfd0e05989cc92ae3d9b6ae1087ca250f", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "19477a097f576cc3ae920cf8be1052f95511c5311e00560580aafcf5c3db93bb", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b7e70c83747125bd1899c8e2e1b4bef3ad1f28622b7aa2d1abb76d362d0771b8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "380468e42957e0d6c01f5fcfbf6d3278645f7d6a4d88fbceed5bc0e71c67ceb9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "e47c1ee4db1b4b5e081759eb18a4487efee11a2fd249e939234665df196b7fbc", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "5fffe2747c0ee9fe2fbc651c23b73146b864454838c892120400e967eac27c97", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "201d9d24aba78791f3b9a07735311004488bb02bfacd2e473466b10eb0e0a844", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "168fccaa2611b0c64a3cb0bd8b86df0a2c3be5943c4264e54c273e1428c060de", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d327ba3c5feb9f13ec2aa8680dd1474fb3983f50884d306bcb4d23eefccb4a78", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b040223ef9e9102891c48dc599ea041ab38af3b29fe90ae30f008be4104c7f9b", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "699fc7d6d0c41c0aa73c337bf944d825b7e4a5ad9f633a0d94c3747ccea40686", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "afcbe3d7e5f5d8ca5f09b9413efb8b9fc8cf7f989bcf7a77e974c7728ae39d82", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "eb5ba9bae2edd0b72a858b3bdfa118b44f7ae2c1a5af5f38e0660784639ede01", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "ab4771f27c83a862a86afeb04b97958349fc4fd4936d80d05030bffa03f454fc", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "4fa392229f053c5a7b6785e11269c1c411603f7d45431a621920e66db885fdf8", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "e557f9fd721bddf4db4fe79d398d51da81e28f45c94e647f18e98ddba3d86707", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ca4df9b635118343c431c15a155536a2019d5a802ef4200ae1d091f2f471cd57", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "75b5501007b3b5718cfb4dad5736a110db9f52bede0209cf610f27766193a27c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "ef32e36e48c4601596f98bf1cbd55758ae6792c7597fa20ffb0d8e51c2d1f4c1", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "29f974c5beef050cdf9fc92129e960ae3e2982ca2d853cbc0906635b767f136f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5cf7f7bfa9e915ad1fef7e8f9c7b44ded965a99d99aa324e641bb5d1ec87726a", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "9936ddd97eb842bf942b5ef2b6723a63717de942c630678c0f1da7b8b25310e7", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "89d81d9f7d275caf28aaf673ddcd46f7ac9a7cf6dfe19b1cffec707c2c1850db", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "f25aad6c134ebbeff1328f69ff502b0b36b6800ce9f70794d7a8e7d73545ca34", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "3622c7df4883478f92432cd919604e1045079ba764acd07e9c5c29c5a5e5e3c2", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "8ab6b0ed14e8be35cbd084ad6b3c77f07ca503fde2f49c4cafb718404448e372", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "58e96a523de0034c6415a89c70d699cf7680ead7cdd7f097f6cb10afe527c37b", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "2b77328e38dc81867c96269d8aee9cdf905705fbd03df0fd372e649f35ee8fdd", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "19b4bc7a5eef3188d142c97d432ec03c207d95b6a4d6984f89bf830c8b1bc11e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "786e651c263f149bd5b95f9065bf314d63b8e0a9fe3ce71bb40def5c1d8432a0", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "49c2becbd32d490db8e19e233d891bffa6f09e4d5c7e1af0cd1807cc59ac0944", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "0ebfdface355c6823c3d45ad19f597d81ae6a4e92b28b15f382eef588d0a44b0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "d2fa777d5625aaf4a54a350f1f6d4c4a933e36b874fda9035f5bc780e353ec07", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "f590c34eb9be9fd909b822402254729be77e67dd2502dfd8bdd7e86936503467", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "308bd0a0e770ae416dfe2ba1ac4224d50be36a5aeaf9fd494df009e5418b70a8", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "f56de33939583bdb9be12b6e013cc47b995510f304a10b6ecb3377418f5e7630", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "32fb1f9211419bf24126b5cb8babd65d6aed4c12226d3dd661573ba3f40d93be", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "de340f5f941cb71c4225abaaebd7233373cdb09b6740b6bd01ded8a21c965128", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f9f559446ee724d8483a1bdae0aff70e21abecaa6c7886e149fce55e7ff7f3f7", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "f8a28855ac107547f3cfd8421970880a9cfa2f985f7269bd0120a8f8c9cdeb7b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "1affd01185af09845ce69d9875ea95cecd4217a5600a13ba8526d1aa5dcc51b7", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "4967952a6bdbc77fdcf214398cae8cf2bcb1e37f3b2503409717662985901998", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "74624f5419aab6f930a7b063d92ca6a0d7882a47573cb188d86e36dc8b547163", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "68f69f951da2073cc648239bff9ce392179cd9f3ec20ea274f09f65571062d27", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "9267dece1a2b6564d85b32a4786f6b05e79b89b267a90fafdcac6f6e46c1537c", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d3bf20f1dbd620b683209de0ffc97f5eef88fdb3ac4dab98c66525d7a9608d50", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "5ad3ff2618fc7488568cdbe455ff2c71578d17b7e6aad0db8c799d080790951f", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "4574666d71bf8329e1f48011f1d0e8c815296d365019dc1c3e8163b14ffad145", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "3be9af3e39c90eef1a83a3c6f23151b6d96392e653ad1a84d290078d3ff918e8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "cd59d475ff114db9f82badda6b88fbe35e7490dd70244145e74e33c751f4912c", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "2e817de25d1b7b6e599b9ffc95a79c22afe26232bf247a88d34c918493af5c0e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "7d2db2476544cef820e99ed585ad2562ff04e3b69f7616afa5a5f861452ef596", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "ce34ed0b8468de250273adff4080abd571b25adbda1aa25642db0f3627912c2f", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "a9fbc12eef26ac635e15a97a1cb78c018e375f9305e1e3d43164fabb269fc515", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "2aedde5d15ee3fedc8144f8116fb2af392a7a84d7eb295678f473a5b2ee1f6e6", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "c29ae36cca19ece592daa5694eb498d4b0f32072a10e43ff558caf11efda0c21", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "3f6c110e2624d0af007d550eed2c227acec1cfde7ac0da3f13cb605f6e184ccf", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "848ab1a0afe46196e39109a1e6356686adefc1dc5d70c63c3f41b36eeaf0017e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e877aa186704cb17f01a906bd3692a4ac8d2e540d9235a7ba899b1baac95516f", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "6f5459d01ec52081159373530283b04ad247b93f73f7f15c3136aa63fad5ed51", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "84306cce2c0ebd322e4062c043f52784174ab6f842e8890c719f75a31c9a6f1d", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "02332a3353ec3def93db8a27b60c7ade36b1869da904f181b995abded244dded", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "29e3ec064f816fae7f45ed72e3602fd115116602bd480510ed55a3815f52d117", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "ea10ac09a27e66b1c4a7db025de1c01ad61d8f4f5fa2320b7a8159695643c6f8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "e5a849c325264757ba8843a8c9b4499d6de8a529755e614d4f3643cdb98413cc", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "af6f21eb7820b38a48c98f3c2e348e6a0ee578badf7d04ccb1f4c858a396364c", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "dc7234da3bec5d37f70386a1c39a7711df68ac420e11f97e7c434dbc8bc31e6a", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "9611635d345d2d8fe2be542826abfa1a40c68b3d5f05240f41223debb9223090", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "8f73471d94e3254d087957094e7b95037053faa42ac88f5c07f145af17efa8ac", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "66db48188cf2a0d6f0cce93f5e84ceb82aa909e09e2e4355203ff9a81dda4301", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d88e9f15ca2dfee7e325834f30b18c9b2f59c4093883f1e93d1312550a855b3f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "0d7b8ff591e21295281a3d0b66f8ae7852dcde02b3dbc32f7e3ec81e967a3e3a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "901eab8c5cb590edb1ed1dc51991a0e2766ff9bb04aeca44095fd8b2ab492b8c", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "32fe2861d139e0f587add76778ef80b09f2246bdb5327f13debab7015fc7da39", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "04c4c38bfc4ca44ae238996a8f14b3a877482fd179bf9ed60b0a2216fc4399c4", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "8a10cd6039092d0f9c90960169ccd989a934f77a6569e14270e206e88b6b1255", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "a9d66d1fb0aad3c6be82aef82473fc205654f4dd37d966011deea9732c7faf1a", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "14d5bc7801b88a9b28b6806c9ca2692b752dda3eb4c918c65635bee3feaf9322", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "bd5a36c852a92c9829ba4286c70b3fab764160c068d2ebe90265c463f0d8f134", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "14086be6b420eecb8ecd99e5e17b32790c342e43ae8b9e885fc605c10866f5d9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "37120a439d0bed0d80351d050d71cc88b5b0ad56cca5d2af91cbbfa8bdd0f77a", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "67d3d37c7fd44f0997e4ea9422cf50af692390e9cbb00ca26b92299830831c57", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2f9343023045e09270073df98488117709cf6c16f9928cc23dc27675cba4c176", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d61a94fe0ebc2936e7ff5b5e8e6dbc19ef4f6a76025b32491e8c62dc482e5375", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "5bfccb297d8e2a37452a1e991d25a1dad8d176f87c37ea1357868ccec40dcab5", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "64e5c6c976d6bc1dcc2a12827af3e1225ae55350d2ac7c280cb3d64056fb9ab4", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "3bd82a84d7be06ab5abd7b9f79a8a40353ee829e448ed05755f23ef833375bb1", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "3f777d2da6bb3d87b927f1bbfaf76f64d603189d67bb851f22e3db28d677aa9c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "b54b7ac40fd643a20b0988ec41c575d59e6e1c9c77d5f128038208d610625f7b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b7e9dde8f0b4db9426c791c37d645d3014117233cfba2298b0ea062f2002e0c2", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "7f8b4bd20f1b9b739551c36cb5168f628af682b65793d0a1a43d080fc1798514", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "a05a303ccb3416d8845c328a124d147cd1bdacf08ec0a604569b8d1816d68b61", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "d1317fdff3ffa3abd7ac95593c31bfe9c037dab5a87b67c29f729f7f9e7c0c41", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "b4ff1235ba6d7e701e4e28aee41b95f02dedc2dbd54e14a0ee2a62720034e52b", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "b53436700307ced25417970b8349b3c43bf1776eeb66e216d8a39be1b1298b54", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "c25fa82a6b86c55927f9c1d999467c119a4617ead4bfaa5814190c780625335b", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "350694de300cc5392a893d42e5669c7350ebf3be3f695a6d6449da2535d147db", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "9de4ebe32ee3d2230524efa121475092c73b5cb4b41b7aa34b834bb4ba9565fa", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "761c4038ba8e36781504b363b778330adaf69574309ac5eed6c27b7767256f25", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "fdc52ebe311af5f2e1f053f6ee65bcf92d993987ccb1ef9360e73ad18d9db7f2", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "bed8f6069d456c2b622fc39dd5c926e35b1ec05e38fcd20b3661ae390bc5c3ac", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f62286f6138a2e9f8841042527a5570a89ed2512d03651eb2878b72fbc83d13b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a3806ea5b3e603d4ec882aaed952f416cd0ecc28e1615ee2ac509617116a39af", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "bbf44f618c9b25065746dcab1fa3d6d42ee7c117182a36d93ab6e5babff50558", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "eccba383653a479582e5f539a57bc1421a83d3e1ec9b3a55d2e5048e69626854", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "8db7c329e2af3f3581c8eef36481dae3427e7924bedd96512dcdbc5fb6981efd", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "97c611aa454c6d11ea921bc1c0eb8208c8ac2501359e48b2150b5b3321b73b77", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b4d38bec8792a7837a0533bbcdd3db2e747ba4dc3ce59c4fcce62f6d1bc9e145", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "8b140bf53a03d08e8fe9b4e482928e4fe8d3bcb30c24043615ebe30bd6044dd2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "828631520940eb2a054ad624be83b32cfe76adfb97e63d883d28edf9e17ed6de", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ffddf3557f742fae7db3ac652fc8225fed06662db5b20f98b4bbf5a1dd80d068", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "c75fddba0390d58b04fe067a2b7846accda6dfa701710e1e7a9b472a1f288de8", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "9e225b7e670b919efd0a0a913353444ec19bf62fe653706e84b636e8e04716d5", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "a25ff767f446db73456661acc6139fef6a30502b524f5cfb02924bef3272fc68" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "6cdaeb41316b86f65320867680e76b1cdd9c4b890e136d90b4c128dfc601dbed", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 28000, + "training_contract_sha256": "e9ea284b76e5ad80bc54e88bebc91fd8112b8fa0a1b1dab774c4e2dd6c2b8718" + } +} diff --git a/models/ppo1-recovery500/0/_CHECKPOINT_METADATA b/models/ppo1-recovery500/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..7a616c52953920aba68017c49ce3e3818566e67e --- /dev/null +++ b/models/ppo1-recovery500/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786225450590947951, "commit_timestamp_nsecs": 1786225451283851473, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-recovery500/0/default/_METADATA b/models/ppo1-recovery500/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-recovery500/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-recovery500/0/default/_sharding b/models/ppo1-recovery500/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-recovery500/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-recovery500/0/default/array_metadatas/process_0 b/models/ppo1-recovery500/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-recovery500/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-recovery500/0/default/d/d9223f38b238fa1d02f4eb96914f5fac b/models/ppo1-recovery500/0/default/d/d9223f38b238fa1d02f4eb96914f5fac new file mode 100644 index 0000000000000000000000000000000000000000..91f32e191f86ab04a011afebe61845d140177b05 --- /dev/null +++ b/models/ppo1-recovery500/0/default/d/d9223f38b238fa1d02f4eb96914f5fac @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5d22e46d003f1d81909623bcbe6c2852093ff8240ee2a2b9d2a8791ecc74a242 +size 112897 diff --git a/models/ppo1-recovery500/0/default/manifest.ocdbt b/models/ppo1-recovery500/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..acbdc3aedc320023f0518ca17f1cdf9dfdb32f17 Binary files /dev/null and b/models/ppo1-recovery500/0/default/manifest.ocdbt differ diff --git a/models/ppo1-recovery500/0/default/ocdbt.process_0/d/0bb4e45cd475cc4e93c21893284ec6c7 b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/0bb4e45cd475cc4e93c21893284ec6c7 new file mode 100644 index 0000000000000000000000000000000000000000..e2a24cdaf59697cc0b1807bf9d5a88094989cd09 Binary files /dev/null and b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/0bb4e45cd475cc4e93c21893284ec6c7 differ diff --git a/models/ppo1-recovery500/0/default/ocdbt.process_0/d/14825823161b2bacd0a1449ca806f4a7 b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/14825823161b2bacd0a1449ca806f4a7 new file mode 100644 index 0000000000000000000000000000000000000000..bc60aa48782f8d81c74718f0cacdd2b208e657ef --- /dev/null +++ b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/14825823161b2bacd0a1449ca806f4a7 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8ab8d8146a9820e5b2ef7e1791437d241aaffe4d68d85c398080f0ff4bf3d3d +size 2236142 diff --git a/models/ppo1-recovery500/0/default/ocdbt.process_0/d/b49d2426d299be18a6217211a14c78ad b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/b49d2426d299be18a6217211a14c78ad new file mode 100644 index 0000000000000000000000000000000000000000..bc0994fba96d352225fd79104ed55c39844fc5db --- /dev/null +++ b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/b49d2426d299be18a6217211a14c78ad @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7100554192b36cce3fcd80399dbe3af0265a59610943674364aa1c3863169b0d +size 62070784 diff --git a/models/ppo1-recovery500/0/default/ocdbt.process_0/d/eb0be67ba49e905ab08a513dad7d0573 b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/eb0be67ba49e905ab08a513dad7d0573 new file mode 100644 index 0000000000000000000000000000000000000000..6439a91f1c492f01bf6c28564de50e9df87aab4f --- /dev/null +++ b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/eb0be67ba49e905ab08a513dad7d0573 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3773301888aee857dc6c2152258cdc0c47b582f8e9d49c29430dd28540d31ba3 +size 515039232 diff --git a/models/ppo1-recovery500/0/default/ocdbt.process_0/d/f1f019f57442b0d4229381f27b4c5f6f b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/f1f019f57442b0d4229381f27b4c5f6f new file mode 100644 index 0000000000000000000000000000000000000000..093b5934caf31def1471feb25d4984ca97685983 Binary files /dev/null and b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/f1f019f57442b0d4229381f27b4c5f6f differ diff --git a/models/ppo1-recovery500/0/default/ocdbt.process_0/d/fa2d110d5b9976dc317c2e9eceab8f09 b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/fa2d110d5b9976dc317c2e9eceab8f09 new file mode 100644 index 0000000000000000000000000000000000000000..61f0ef4b1a1a0c6fa368f423e3f048b9ba69fe33 Binary files /dev/null and b/models/ppo1-recovery500/0/default/ocdbt.process_0/d/fa2d110d5b9976dc317c2e9eceab8f09 differ diff --git a/models/ppo1-recovery500/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-recovery500/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..58a35c5ef5585ba7f4731b99f13e48ee38807f9e Binary files /dev/null and b/models/ppo1-recovery500/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-recovery500/checkpoint_manifest.json b/models/ppo1-recovery500/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..056b9812a1924095fc514c9a8ba420430366f802 --- /dev/null +++ b/models/ppo1-recovery500/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 27000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "e0cf2e33b48230476de493d4689923ddee970efa0f45e13093fbefb98dd672d0", + "runtime_config_sha256": "e0cf2e33b48230476de493d4689923ddee970efa0f45e13093fbefb98dd672d0", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "e0cf2e33b48230476de493d4689923ddee970efa0f45e13093fbefb98dd672d0", + "schema_version": 1, + "training_contract_sha256": "0a97fcc21fb7d5f5f51b9d3f915228b965526d3e31eac9955a6a46cb9c8fe9cd" + }, + "sha256": "f78377e25be03f481a94977e18e93e6c1116456e90cdf3c63d5bd2b82a843475" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "a16571ea96e459bf857ff49f78328689afe1f754ccfb911fc4f0bdb046ce1f6b", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "af33adffd7ee3b8949e22f4f2df346f5797aea400cc7e1fa419aedf6d021a864", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "9aa689f6eb193240cfe3638ac0e514e782c42787ccba9b16319bad6e0e8fb3e0", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "c3b2f2196e42f9626fee241b8aa53b21ecfb561df924e75bfc35fc37268dd42a", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "5b9f4e5e4e32e626e7fc97982dac0fb0de1661ea87d616f7b0d74ae21c48d67e", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "96fbcbe16c18876ff66899b8ad298c13db56ab1aa6ea52bd93279f283e967187", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "c6d7865dd9c3d584040deb5cea299be5c54a877bde5d0080046b369d5520b6c1", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "7ba901eefb288d8f0ccb25926490b1a69a76fb8f548d5a5b0dc44472953ae1bf", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "085950d859d722fa6ebfc5c2d205fbdeba852a6f78ce8c9379df7593f58c3837", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "17c7deff0f35b2a46580a251edc0c463694a0c26b62bead5be7df2afa188fc99", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "5a861216ef496a5e0c48c23eef7d021b2ddc732a6d88b07593200f970e79890f", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "060e815b588c0f373a49cc0d4ba915060f858b3441f68a90479c9d80729dac08", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "ed8e555805516d3f909eb4edf5dd516ece816203d81ed0f4072f51ed151c0396", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "006dcb6796a7149f3b1e3dbd6b05390619d2c7a514758e0eeeab21f6b2305b2c", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "4fda5a1c73aa8b9d1a146ada54b4a8e27727816f4f97fa6ccd32e28750da26af", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "f3ddf70150a4a0741788f92000a21249c2f8e0372e77bb88120a9c3d8ed64e66", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "e6771923ff20783e099fb1d7be3a834701809b65f16cd4c7c6658e67d1400648", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "093062605b277e069179e72446f44106fc5c59fb285c6d51428752f74519ce29", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "fd445f8161ab08561aa3eea4f58f6c8c60700c32e4d926cd00967cbaa67fada8", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "33c1e46a6fa874689e347141691d0ffcf8401a81c769ef7ce07c6368be40df1f", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "72bbd0189bcee56bedc8e173b909053a11e64f542acec0a13f051e53ffab6966", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "eccafb0156145ef68285e4cf9e622d932b9b16532820f219f8290194f187cf68", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "217934b28e8907561f1bd720e514b8a705d5c522ba14e66e427a97a41eb3c234", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d4a542ceba78c36631652d31a093f00ecd1501b59cfcb0019da3392cb2a84d37", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "4e5e8273d6d6fe62b2d316aea8f36f38d0dd2babcc7b69de506b295a300e5f05", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "def241256efe49429984017ff87a7e3b8b810681c46877edb001950fbaa842af", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "5647aebc0ba7d6c583e567de5502bdc29a51a0bfbbcddb534a7b5bf56440c9a5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "237eada66d6768b9d7ba8765f7c54c37becad18bb61e291581251a55bc644972", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "8fb99c5f11966664af5a075d7e42306c8aae58272c64774b7fde6f17e1107ea2", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "01f611bac4a7a960c14cd2018cbb3ca0da40dbac909712e47337f3d34c474298", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "60a5e4036b4ebc031899b4a27f81347befbf2e2a9f1482afb43f15bdd74c4704", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "21148b183e99d0e3844877d25e0e73f6ab4929e90c380a3dc37f2c90c78e9d65", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "dd2b1067dd4fc09a332cf19f59594170ebaaf095e22399e70d34dfe017ea29cd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "a638c74240e851abf7d43a30cd0887cc27ae6d6fa7487228b0ce0ac4fe7d7295", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "733fb176ca517bf35840fbaad07bc3ea70359926594f9ee08ce0686a2cd833e7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "001b394646995bee15cb28269130151762418451aeb59d041380dc86a33bd018", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "365b827b13178fc9f4fc57680030fb547c2422332e2bb44ca8453f23b369333e", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "4caa870dc963aac8c16dd0770c828f296f2d33a6115e7c4f2ff4eec5f850b61f", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "7748e9e8a402fd4a93c92f096fdc649cc3e23ecc04bdae7a95ced8e3a4cac460", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "8361981fa82f35cd83c0179014b7ec7373206c8264f6068345977d6510093d55", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "c3e318e2b472fe8f4c09b59802a8d6ff6ba82cb77178371a73c5ee38502576ce", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d50fec747be7cf4bca359b3234ebf770f4d22b16d6bde41fe3cc5862b17784a0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "91e816d9986018ae6b66c30cb9e8b5da78b5c8beff82fba6e4b5b9696a9914d1", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "71f7448a8c5ceadb141c58e0c9e27bf3cbc0563a4587bd5da9354253d0fe19ae", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "c18ef44ece6a6a59eeb2612ab655ca095f8299e6fe81b42549ab0daa2cabe466", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f60e2b9d4ee3e68f7930290fa081086551ae7b82e61cb49b8981551503274e5d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "913b0a6b84ed14c6948c3765a075f06204e3edf3d60698b58c32178eb5d75d3b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "dc3fa82d2ac1a9590acd7961cf97d66171f1541c6d3820efb02443dc81e781cc", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ac573fb1e3d4614a09a721ad47fe5b1ae0318c54e6b718912034108f3ff0fa39", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "684da1ff61bbb3a2dfbd9334b5adbc6075d3085a325b3d96ffc01b13c01c0172", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "eb8224b2baa78d84d63550c62e7d42bce60719ad07d8502483f73998d26f8177", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0fb4eb979bf82cdc2d482c019bd34468468e80a85660f1f08c2be141ea8284f5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "2fa81c3da81b90c2519b8563f9640eee52627947bdcab615e4682e6519efc1a0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "aa69299cbb734887d5728d4c3579e2a5bf6aa024936e0357aee23cda829b9dea", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "13c4b3d0ed266d854e9ae694b9a341f7f5bdc8122114d61c5c68e9c036150de3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "c930dccd9125f8f1ac7fd540d802ad6c18e7837df66445a663e9e3505dbd12ff", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "cb6db932d598160677b8eba389e3536c15dda1bdbc2d019844aa5c12aa6166fb", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "da0ae0f412e6c60baa34ae64c57b71edab7d5bade0c23b6673ed58716d703db3", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "fc9f81db93c67a117b7185cb2cec5f847c8a58c664469cff2912d78123860fa5", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "41a97d7cc9aee7cc41f2f66a2e40ad41e2333e3077451da1a9be1920a65a1f38", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "fd2ddb824ea7027aaaaf582b620f6066d336d092e696dab2b162b56c1d476891", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "9075f39c1d6f87e72bd0bae0d725010cf7d9065cb87fccdfb501ed163d48d54c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e8f34f869e66cc6be9b358e31148eb41ef9de71f94edcf610ed166a9fc69f8a3", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "2195105304776a0a239d5387e199ae84a69edf67d3ea128bfd7d21312ab85f87", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "02f89fa09dac185ae080720b976e410dd6551a9228977020f5dfc6f63a5c4c30", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "44792a6213ae7b80b30a5110266cc39f0adbf43d2d3eebbe50a8c17700521718", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "0205e43f2fdcde5633bcf96a909c66311b1ab178c541519bbb1c5da6a0f3b405", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b6faa1ab7ea0dc78e3270856f719a5f42ef19898dfa2179409cb92dad0f67529", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "9c940281f987caf377a270978ba5c0ac07c571d0c40450942eb181723842f46e", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "7b44e2971ae2543e604ec21bdb839ff4a0da9795488fd96e41c59353829511d9", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "44594fe18ba92944b48044ca08c8c44a2664c0af56c55acf13f48bbffd32c877", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "dcdf1b312b630ec869a7b6033f5c3e84623b90de2ffb8a44c3bb041fb6e7ac09", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "9334159f60a377a37590d26d8537ad2a1eb92a9b82ec108aa19dc66f56b83f86", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "c68842f3877bd23c6eab5bebb2f1747553d63501790ee446fe5f858af5c969f0", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c8b4457765b0ef1917b0b3e611a0bb2113b9e2ec6409d9275661bed70b38eb27", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "f0c7ad74ebfc3d2092879c84de15e55fc536b56a631bd020ad1818401d3e63c4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "c0c8c344bed211b68453c9068eb0ec2364bb1c82c27f08dbf0597a82df01e6ca", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "7e719d2144e37ae7095f050134159930696d06ac45c68c924598cf09fdd50f55", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "aaa87b7c0e344aafe3788c3a7475472e478f1d4d770df47a49b5ed86bb5e0d9f", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "3e1dad9d71c1620f17dc541806472adf8fc20adba095acf6fa3596d9f9ba9471", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f6a8d1269b3974f0e215c77bac7e6bd1b5d37c1278d60b0dd66a06a338afe5c0", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "70f487dbb467b7bc71daeae67897fc8004142346b72e9810b17ca05c1a25b119", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "231cd46bd4d91d99df869c2b9e28a6969b9f8af6d4ce5d31f69f6a3253ee0c00", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "694adbe8b8594514e5f59492a99f399dc1a097ed2f1fc0dc135dfb2b2254a77d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "28da2b108ab7eb3f633d6444c358cb99131dfc3f9fd558d6aaf8dfefc8d97315", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f9c47cba463526bfd83e79767a74a76fa562f1da637a200d225939e649c00bd8", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "fe7b337d4e632f13253707f2ae68662f1dfd0221f6145e5648967cde87afb840", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1c0e64bc456ed8bc9be73d50f49b0bf851aef531d8a4116f2f8dbc79e6aabd8b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "0514f24b1a48ac8dbfb65efe919f892d2b38162471292e724b234db95600b36f", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e7f71b5d7629603e7d0cc69a9920b1bb7430c81001004e9a1ce7d37014859acd", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "6d40d1d12ddff4cfcb67fdd79c543e1ed3df0f5a99e5fb62b34bd1eaf122ab45", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "205105d711cbf6a477456881419048b735d3b8f1a3bd43e5ea6853e0d4d093a6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "212eb51f04463543c69e3ed8a751851c198aa4b06946ed04179ae8b9b3c76718", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "3753b58702e4c8d20b8b98e2710e7987980f2d0e44e136e30a284dc0d3044250", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "9d1ad55db913d501ab68d9a836b719bcea3f67d58eddc8d5775e4196f914ae1a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "6ff2a4989a9679be4a47919c4a14f98db380ffe66038dcf67c1f6a93262055c6", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "4611b53f59b784d088f613902fadd97c44ffef66891b7ab63aac9d4c9d504722", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "ce4495478eb4258e7dc6d67912efdc0eaa088beac39f9a144f43f4f312b13f4a", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "9f8db0c6faf5fc135a2f176c6e3207dd2bb50a2b3b6d66bd21db35c9fc452226", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "5b3382ec5dfce31a69dddb4bee692d768992de9f857fa79f78bbc8010d0eca16", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "83fae6064dcee8ee17640908ed9526c361e01374fd899ee9029148bfaa06ca4c", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "ac5074e179098a75b3bde31eaef7f5a1aec56da300d8d86fdba15b110319c6da", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f2018843d608d1091344e7e2df67230c5cf202f3273e1b9a5adc86f77746fddf", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "dd1673fb7a6c776113296b8f3435e415aba6bbc8c323d5a7020705bd7daee40c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "78e27eda3fdfb7e9431761a9ffdead900cfbd4242e972b9f68860fe76f50f7f9", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "83faa50513ec084f3f8f0fcb2faf95853fd224c03477b2b9100d0bd5f310516b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "c19842bf669bb321ac625212b924062b31f6cece44c85801829922679525cb21", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "c5a2eae73cb482e9fac18632f75e0289e205828f2f77844ac4df786642cf7577", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "25202c58ffacb4548df4939c326416d852fb58e135cff662f8a819d065b67e1b", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e278de10e6962944fe8bd4609d1bef6eaf4675ce38ad353ae6f45f34467cf5b8", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "f65fd52a6408490be678791ea4bfdc3f2d2f70097e80832b7fb95780503d73c0", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "c5f9b304a911f4f54ae2499f72376f53af7e571e257b0a03d2a73987ca9314df", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "f085e4de837b7000a0c3fbf33b93bf530813d0d8ec57c54d961acc943b167f39", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "308f38a6f4bba86f56d9a24b9ec826dd8f236dd34154cfd1fba57b9acb5ef76f", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d025e3e88a1947f0b081d24da7a348ae863f858b21fd85c33306953033278eb9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "a1a813d77ebb66ff79ec03e94dca57f2cb6d942d09e563ec88add38e9e958d4e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "eb738fb77d24c536bdb08c1b1f0f0783fd35d723f76cca71874ee3d64d854e84", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "e1a38ab1adafa9a4e0ac5af05cf233dc97979d91763d97cc02ee4f9d95a41c59", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "38f72ce6b2b5a1055a42fd9ca4c551f4b17d31b4970e8a958a4f017e1fe3aade", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "4219accd914bafc8d2bca437145ded1d149f0a2e5ec11e8100dfc46f1c9d0cc2", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f4881f60721140ee1618de3b69b9a28a363d2febf96c5cb290717a0d0ba1f02d", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f98ebdf10b05f0392ca7ccae6e7886cdd77a0e69a7f6a557bb4ca9b073c4edef", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "df2486779b7986adc3d87701d0ff4eec2ed6472bec1d1bae38d1b81b37c36cae", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "56222de8ef900bde67b3342c248cde33250a008bf6a1f2e2c48ca161202824dd", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "463fc5e4e4115b64f624ef5d4a5e3429855a77720314f3e837528c6e1f005c5d", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "76034c3e955e7f3c7c88791bb5a31cb4c431792084742d503c19c8f8edf7bfdb", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "f9437ce06537587ed5ac1a2e5603bd75c4a7d7b7b006ebf51998863a33d20173", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1c580587a1611f4400cb703efe7f7f3f49ec151d57d436cc475a70b01c4cfc4a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "206d388037d99f3c66528c94006c4b0a528ab715e8027c029da094cc36e27a2e", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "4de8b277366c987e967bffe38ce89f67dc444a1fa822a2d697be346ff7ee543a", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "de7df2692f84d53ae9bbb6414fa07ac14111405c31044573381b1bb5c8cd2c76", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "74febe7eaa6a49118f9bd1e6905d0dadcb8f45b72957b049e895abe7124702d5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "0ab666c09d9f9e6eb653c910e02be7f91d4bff1b101077067335a62d3906e4aa", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "854f9421f6eb583565b18f7179759920998c67783b4790319074ccebda8d7154", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "0dea152638d97368964fb0f27db06b9004fe52372dd74253457e562801b11566", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "8cacc52710dc2fd2d132c89d38928f615ab18f261598d5c70556723c6a4ea75e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "f7f6ee6017359ae6c7989ebecfcd4a29e66fa11f0bca110df9ee1d60ebb963af", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "1c47e01df8738e46582572a5cae2b6befb2e359788d0315ec329c802630f8248", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "e5940f2f626b4d5e187053847c7317362237b15a289614d2d818fbced5e8f691", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "8b38561362e3212252caac663a28998946f7c44da9a1d105db4d5490561c4c31", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "bfbe2a2dbc02f530efa866d8320e6ed6a8ee14f34ac6a5a8c1ccda9b36ef7e4d", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "82a32480f86b06d2cb9b14c3eec3eeaf1e999e188a7d778af97cfffc1cfea006", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "fabe349852e3b3031967746b0d22c4503d96f623515e38e9d9486ee930136628", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "69688a29c9afefcd10c6a90c50822f1a9f22e1537709c3e2e28aceba3bafb0c1", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "be6235a666260dd811e42e2c61f6b853dc3e5a268fa24c824cce013ac042c5cc", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "7e260c9430f055f9a1defcaaa02bae95fe259d265a73c5340be8bb8535406846", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2f3cf23c9da4b5dd721ba46ebe29c81fe50436139c7cdf4ae2352cfcdfb777f5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "f5584b1bf58db6aa85a6ba334f336197a69cd0d5913c955ec8c472adae794229", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "5b497aa0b2b70dd962c2f9732d62cfe3c172b5749be05e29425ba9ba82406f75", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "2dec328ced1dc514995d0d8797a364c8acf4356a89897b9cd6a3e4cff5eabd6a", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "b0f1c497df12d3621dbcf45621f214cfd78872e1ea782320bafb58c30672d1f2", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "e1655d221f898bfa74aa485ffe68bf278e816511f2ce2dc44ca5f016615adf2b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "649ea4a984a11f6ecff479b57ef6b3504aee33149eeba623e7157335759f1628", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "4ca9d4630390c87b0b08533a5a6ee5af24b1c3d24e20d0ca8a2df4b18be7573d", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "e46d57a201b63c594b670a0fb28c864dfcd6cd4d3932e7d7ef35d893da16b5c0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "7c283235fda1e8bf76b15a4f37347e7d255e8a3cb651190746c5192193d6b4ca", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "522f449d7d7fd892e3b364c2e56d94bd2c129a2dbcb1fa9fff61d0e408823c6f", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "9886ee0334152f46134a5de79401291365a4ae102e91bfdc7fe1bc6d7f44e3f9", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "9661670854afc0639c81d6fefe52bed298dac8e15eb19063eceaad39cc752ff5", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "56d92e87e10cea2c6dd19eb597dfcbefcac92e655ca8e2f66e8da590ee979f95", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "58e66be40746204bd17ade2502b43b74b258a131f8592de9109af6830f83931c", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "be5d1e7835a1f07a1094ee1c197e1ba799ce53c55b5c96ba90240774f84627f4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "76d3ef03a31ab0ada060ac8ecb720b898031e118658ecaf8b0b40252fd32f9a3", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c1e03d21c739108e10900144c15e9272a4d53b761a2d0fd6f4ad7ceec907ad42", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "ebd2c1ce09d279311bf084876e353995e7c1d9d4d3752e54d00d363474f03687", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f2c90e1afef3d3189fe01274cec6ef54bb8034a6c47808330fe7b27332ef1c43", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2956cd92ada3c005a66557cb8fa91b52fdd1ce5c2212538703b3a9777131daec", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "faa5d67ae55e227395777f87473c00b1157dc2b7a48bf877f743d88b6335326e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c7eaa590596f48f4a65a3baf90a2806befae3d97a2e300f94712bc09a5cf8dce", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "4ec1cf8f16e6b76c539564ef019c0ef3fab6b23f224bdccb91c0ad9cce50335b", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "14f09f9afd7684c7ecd17fbde78c51195820729672213012d16ab7950ecec882", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "d0d0b9314ec43c234f94234eb89d46b0f874e8e3e8fbc5d24fbddc323af1ea34", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "fb551200196a63279dc9cb67e562588becd9d2b01850ba9e00e7e46bcce4c1f4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b9103b7077e0b7ba8025b9844293cc748102b6aa7bd3872a5442ce0bec127092", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "a42f17fbf9f1b832cd6cc87507edcc6647b0e89b07584c823d42073cc8899b8c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "b793b05aac8b4a16e6f33d328db5fa1fd713497d2bd2ed2633ec32c265322c49", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "cf8b030e59324f4e818a1ce41c4bf67d84416bcffcccf2ed1d8515af2d25b000", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "6e9ec7e4639fc2702fceebc7fe0b785f6bab47a8f888775d85d51a73fee1effc", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "65ea6c150740012cb8acebb2667daf9f9ff7624fe9d7f1bf7fce156c415b467f", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "c84dba5264c07b07a2dc5fdb34a7a6e86caf0dd77327cca701c6e9c0f2a0973b", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "906dfc46ecc42315ae79352b0341b5402a6ac7fe38ddb5f98d7974ee50540a8d", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8c6617056536c2088c72eff29f2b07b6a651cd05c6c59f91c09c455dbceef00c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "713850bf4c50ba81e5b876447ac3dd12930277423f9235116825214e9f17fd88", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "cef15bc98d9a8e25e0aa7189ead8cd1dbbce8b72349067def6ae82a777bd977f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "923377805d2db012e208240af576a598411fec6c03c0f68a775c689ef57075b6", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "fd25309d15f42a54558c3cda2b38c881cc5bfc95df2710f73d806050dbfd130d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "da3c5164fe0d1067272859a487feae2a28953c3b688046dc1088346faf3d659a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "55c6b82458904aa4cef26e448f98c44a78b72c37a3ffade373bd733249ce5aea", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "a501ecea529c96f24a24445dae4ffdcc6d168bf10379351961d8c71254209f42", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e4d0e1a0950fa40a19eb794225a6b85d3dca1447076e08d16808525b80f80d36", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "0e07f3240cdc6215cb1ec21ab96b882f9716e07c5cde71b4819c1ed12a6aa20f", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "6424e047e9de2714f8cc3f46aaef71edb57dd09cf7c490261291acfe1058df56", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "c44d765f9edad929879b9f1b0433e09509ef4353cff2b9a90960518b76b36d38", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "2a435f4aff656d1eac38a789975c604fd8b08b7e698a021d9ee8b18c49363110", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "f1581cd306a683054ae53529408411827b03d0d9fade7eee13fa0890d9930da1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "321ed549ea059ec1dd2da1d35a648436ff52747882e74472d1356ae361a0b088", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "b4a1a04d69edd6990016b35760850cf3289dc18a940bda36bcb2a635ef09e4a4", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "10f9f4fceb7c80b361a2bac6419541e514689663e60bd1eb6065074eefc18f26", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "14dbe9f4d2d5058c66eec18abec9ce15a45a6aa5e8cc62432f2e4eebb0b438e3", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "80fc60ae3f75ebef512be401b7c57f03bc0495783d977e5d7e93c69bfca5c16c", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "443a154f1e7b8bc451d1df5d264e665b30d83e854f19b8d26f3541dc24c00aca", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6e9420f24c4a4343dfe92b0b15d8d354e099da522be290120a763c0c12724afb", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e00250ffffe4106a9edee7e3368fc6fcf140ff8432a2f8b645490e8fbf7dcdf0", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "f89b87ac13cd620f5e0cc64784cecb364c36b64ff881340cca52568a0faf31e0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "113ab0969b38e40218fc11f675af82178bfd1558c4e48602bb4b5b9f8cbe440b", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "187133cb4b283c7403647bcb1d345b9d4f58faf962e64fd7ebbf345b1b013260", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "04fea16a0f65e2f18165281ae8a842ed9930bae5a9bbeb42d9c525ed2770d560", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e9f5f4969db30592ac866562c04d824861773e5d287215cc96a3db572bb2c510", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "042adb40d7c6ceeb25d9bac0f5806d563cb360f84b7313daeb3940499e573306", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "5287b1766fdd2db1f518f87b77c3710aeb0316176c54d4040f7bf137006c7362", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "3417d3004d8e14c6eb5663a235aa5dd3c73daed707a3fba5df4aabefb6880f7e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "630356a33e365e021e8cdfa9b7deb725e47b12697c1017da5da868c019b5eb49", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "1a9087ae25494d75f4729d265ff7946a7e4ef07c287b0d7923d8d7ec1a86e2d0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "2faa868606b5f6079bdb2fcb2c62674389bbbfea26d612046fd76cba5e7b5c72", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "cc1290fc8c8612208d8de6fe394b781816cfe04e5e06cd3288e7317b2e380a44", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "6a8aa6ed670d2d5df8d9424e957ab8a253e97587068dccda6c7186aafd742a86", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "30fe3088de752b9d1524318176bb716193ce139affba04756e782f4ac845241e", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "c3c8769c1fa995ab670fe1048ba7b9385673f746cf241cba92ceb934f0ffdb26", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "52c0c7bffcd9c7504835fdc5a9dd413b84661ea3299ab8b7dae7bfb2ed1a185b", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "03b313d676858d02fa7e8777eb5d6e3c66c21a863225368ab19eaf6bcd44e939", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "7fd535a373a13951fae322025b2b1c5f05ab9767117c51191849a92c005724ce", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "a5ace53a76f16f99264933b477585a9eaa6f40391cedf119187c675c5fa29d93", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "c2dd5b2fd1c687991c4613b7afdea6c9f0cb57a9836448ccb7260f04c2e6333a", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d61b4b8516a18021b88ac319d975ec9d2830101facba9931dd3d566be0fda782", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "6a2196a5b83ed0ec8fdf6b5be8a2339cdddab89a5d8a77618f5d6c18a8cb5a62", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "7b292ec2c04e04a311e8383879a1e8b431d9d5ce897cc3e0369e7d5c038488ae", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "e04e66ce91fd19731c36b3c5b50214fe259ffb9f9f59cb29f3be7f74f52a5c1b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "dbac5f72d451b02799517c9d10c87c863eb05ffa9f2b46a56a0f86de92beb54e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d72f51908e0fc6160a0d86818218da952236b7aa20e40023acf4ec84753de7b2", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "670f3bd1b2e876892f9939481824ddab8399fcb3be754a3d137108958710863a", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "84446654a09d45142ff01db0e26e6a44c13daa48a09171d82da1f063674a5834", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "8d99d5127ff00e28ff1b93c898e43dd19e44b75825bbd973990b6dd3130a9e44", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ed061f2494ac329d694fcadefec688fcd2795ee64873d7e46825bd5cc1596c56", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "d3d63cbfefee712c5b3589d752934905c86b0c87eddad65c2dcd4e3b51606d2a", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d1530df372aa9294d6d83f6932039786866c16417b139c12b50e00dc2cf6350c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "cddff31a92f2324f246fc4c5ebeac10a362f60aca0e6b29823477bcab975323c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "84c6d2e107416a0e4d5d97e693906f7fb6ad8ceb8e6698da2a084625a1ddcc80", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "b9a30003e467463c0f217388119024440bedcad60ea410923546218a335ad3ce", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "ac44fedd87f28ec09ce6eb0081c6c126025efc2118c489d03095c8ec0e53bab9", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "92afc5882079f5fcc44c87bb0cc702cc1e39d2e048e371a8f2f987174579004f", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "c56afce675af14c4b46d5b2ee0ed49fdfd2159746bb8cb29e9243868dd89f085", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b2f7fc2711e58e656100c0ac7dee868238628ca689c910498f448c4e1881f6e5", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "76fe5f65825f3d381f0cdf19d5d317b72585e39432b22c53bdf05195ddbb5a4b", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "fe07f71851c6efae2c3b2f4558e3d396d8c6875dbac187503331d74d4b4d946a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "8374618fc0ad545ee64192f473899478619b5f3c3cf5a10b493bb1b3fa7427a4", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "b34f52692e7c8d33b7cab71a48f58006bcc7d8bfa8b8a67fb868da58dfb0eff5", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a1d9244c8cfa9b922447ed8c5b6463871715ffaed6586944461cbbe0c7bdcd7d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "3e4efc02180138acdcbfd56e3bd2cb2e7a8beac7fde695d417d7f8c6879bb546", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "3048f75a49c93f07882bba6b90c7453bd4c68df6d4eb48f1e67fe98e792b3f7b", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "1cc7a75253564625fd7de15e2459c1a096c678c79440d37b115a15cae53c3390", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "c209514a27dd7f7f05dd20f99b29a22559a5d51d19aa2ca3e58a6100f328051e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "2711ce55e5a428a2eb0074429fa80a2e22b724a669cc3fd311259af4d23a909e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "b5fd95c5fe7d3b1a6f3728785b067277a9a57d5768c10f4297143085b9277e4d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b157cb4c408f8056490a4b33dac35c7069d5caefe5df5b41a7ee84b1d337d115", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "342ba7b642c70365c946d75c4ae444bfec14267a6af9b116e730c6aa799ffc62", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "1de4e835027b1096331f1972a20aa3bca4bb4fec4bdcbac06b8ca115257cec36", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "8209ee7f801d9f48df018b03e416289377a65e65cc1918e30701d455f09c5902", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "c0d80ece428a6e1f02438cc9ce1b88dc427e68e01636bd663350890f7e013ba7" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "79cb20b2b8471aec40dbba6c20e51799c3de27ac4affd6aa2a7a041a900f1100", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 27000, + "training_contract_sha256": "0a97fcc21fb7d5f5f51b9d3f915228b965526d3e31eac9955a6a46cb9c8fe9cd" + } +} diff --git a/models/ppo1-uniform-random500/0/_CHECKPOINT_METADATA b/models/ppo1-uniform-random500/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..90cf33047cd0a8f0602508c077c6178f03805948 --- /dev/null +++ b/models/ppo1-uniform-random500/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786225531640693507, "commit_timestamp_nsecs": 1786225532320664221, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1-uniform-random500/0/default/_METADATA b/models/ppo1-uniform-random500/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1-uniform-random500/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1-uniform-random500/0/default/_sharding b/models/ppo1-uniform-random500/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1-uniform-random500/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1-uniform-random500/0/default/array_metadatas/process_0 b/models/ppo1-uniform-random500/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1-uniform-random500/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1-uniform-random500/0/default/d/fff057e04911c88885541c50ff03e80f b/models/ppo1-uniform-random500/0/default/d/fff057e04911c88885541c50ff03e80f new file mode 100644 index 0000000000000000000000000000000000000000..5e2e5cda713953a50460bf6d668b3b8d7a585aa3 --- /dev/null +++ b/models/ppo1-uniform-random500/0/default/d/fff057e04911c88885541c50ff03e80f @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:467fd9004504edf6b442cba71db45e5cd8afa4b7bdcdb4138ef901e06af9ccab +size 112872 diff --git a/models/ppo1-uniform-random500/0/default/manifest.ocdbt b/models/ppo1-uniform-random500/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..cfc96f119f795db6ec43f80d63d149484a7192d9 Binary files /dev/null and b/models/ppo1-uniform-random500/0/default/manifest.ocdbt differ diff --git a/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/10b5e13311f1fe474a649fff5059396e b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/10b5e13311f1fe474a649fff5059396e new file mode 100644 index 0000000000000000000000000000000000000000..68081f17713bbe397b783de9a378f6f50ae41bc8 --- /dev/null +++ b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/10b5e13311f1fe474a649fff5059396e @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:42a1f11a9a26c89780f5c88c3ddc271669376fff21abcc2d07d52239184b64f8 +size 192405504 diff --git a/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/48a739ed7e454a687ff8054229107ff4 b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/48a739ed7e454a687ff8054229107ff4 new file mode 100644 index 0000000000000000000000000000000000000000..e882d8cf57f39a9d7f8c87d7559f8ad046c7c120 Binary files /dev/null and b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/48a739ed7e454a687ff8054229107ff4 differ diff --git a/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/699c53e18830c1c5729371cfe7992bdf b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/699c53e18830c1c5729371cfe7992bdf new file mode 100644 index 0000000000000000000000000000000000000000..6066ba180d45a534f9a4d508fa15cf85cebd2fe8 Binary files /dev/null and b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/699c53e18830c1c5729371cfe7992bdf differ diff --git a/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/9b39f2f8f2d69fd40c8c5fce5b7ef608 b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/9b39f2f8f2d69fd40c8c5fce5b7ef608 new file mode 100644 index 0000000000000000000000000000000000000000..fc30e0d7a7f4277455768087b0bed3cc813d22e9 Binary files /dev/null and b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/9b39f2f8f2d69fd40c8c5fce5b7ef608 differ diff --git a/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/a721199592d521471bdf14a4837a85a9 b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/a721199592d521471bdf14a4837a85a9 new file mode 100644 index 0000000000000000000000000000000000000000..a1999d8f8b912c41499c27a40768065723181a53 Binary files /dev/null and b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/a721199592d521471bdf14a4837a85a9 differ diff --git a/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/f6f509b14bfe5f51f56c22ccb3910b58 b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/f6f509b14bfe5f51f56c22ccb3910b58 new file mode 100644 index 0000000000000000000000000000000000000000..d254362315ac5dba1210d515725a8d4d86abdde3 --- /dev/null +++ b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/d/f6f509b14bfe5f51f56c22ccb3910b58 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ef4e67bbd789f279e2b611f79d7b907cd50ae01fdc21d1df5973d88943445580 +size 386842624 diff --git a/models/ppo1-uniform-random500/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..495488a45c93c14cfa58ec01a6124559b0c5b8b9 Binary files /dev/null and b/models/ppo1-uniform-random500/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1-uniform-random500/checkpoint_manifest.json b/models/ppo1-uniform-random500/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..b943f734e40e10717e3b50c70209faef8f7bad6b --- /dev/null +++ b/models/ppo1-uniform-random500/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 33000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "676858c9428fe17b0adf4838d1033fca04bc47d6b5a4a23435e71c01570985a9", + "runtime_config_sha256": "676858c9428fe17b0adf4838d1033fca04bc47d6b5a4a23435e71c01570985a9", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "676858c9428fe17b0adf4838d1033fca04bc47d6b5a4a23435e71c01570985a9", + "schema_version": 1, + "training_contract_sha256": "27f686282981093e4bf39aa71163e8792c27b3cd753c22213bb66fed6478235b" + }, + "sha256": "96083dbd4d36f645469cb08c29637bd8d1f52945fcfac20e387695c88dde54e8" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "5debc65090d59f192c90fddfa23550263ae8e8ddb4fba8e7e7e9ef4391b8087f", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "31b19473e035e1000b4a9b509473c62d3f6b4a914ac42a705b58975333ef79a5", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "fde2459aa1c92d520f50e8b629361bfee156abc1f65955cd9e749654a6921046", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "8485e57b4c12f064465bdfc849bc9edc95fc1594c61c44b351d359ca0a7126b7", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "86543cdceb08e0b363c6f2d2e95e0db29d4f4c257d0427c440021c02f237173a", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "8c636ee39bcd40472959aafc9267f5f92ef7ed2af9d63f73ae9044dab14a0ccc", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "41b08807414e168c3547211f428487310f83556001fe89f7f40c25b012cc9f7b", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "b577676d36fbd9892cc89bcba73275fbcf73b8ba32f6dcf177f99e009ba23c1d", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "9a2f3c2b663072bf10a7df79319c7e457fddb426ec9c6d1283779ce9b546459b", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "f415d0541a48500c4a26b102b9ffad02c9a18de934763268b863f4a0e46a8830", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "f043f39b3923127cdbf6436373fd47e30ef91711c8e5b306aed8256bde8c314d", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "acc71b88af560c5a68d1526f35ae48d26f4644bd45034ba409e56c5c91d5fff1", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "603c79a607906569cc34dea4756ee956e1df9408b8b3d9004487ef657056a38a", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "d233c9d5d99a18a3b81b4c34d15165f1b43b1c4d62ef865dcf7eff3b2af724a9", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "0bbf91ddad0e59d55b43419868ecae00711bc5139f18797fbd2b9b558c11af88", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "040e5a0e0afb827a751c286725b46810d7a62afeef124554d845f1abaf0fd89a", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "cdfc385b44e4c4ebf23a846b66933e6cfe9ddd1fc22b211e58903238e858dd1c", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "424857d798846ea96cfbdf51649719b9d48656fd140ca8fea83ededa2f9835f2", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "5bf0c520417b3f981b3e2333eb7aaa1ce90df4681a2f4f37541fb2022290e7af", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "c4a28d4c6ef35bcb4b340315b9d8d9b21d71f7144ea8ad33d2ce4bf43f5d8ebd", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "a7707a12b98c1e1d7840835e5f283fb0b904ba60519f87f0e734d19f40714076", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "cf80b103bbafc5d27415fb71b318f8b65854b2d2aa1bc80a89450918240764d5", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "2cd5b7109dc5987671eb49d8663abf081e004ee222cb13b835eab544072edd51", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "7e00b17adf27de310c79b0777ed39e77b676c34babb33029dde3469d4d06cb21", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "a23b9eff01a5ad16ec636b41792c453bc4d97593f6e1f3b1aeaa579b2c9dce9f", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f64c36724268ba83eebf56b4f7d053e19aed82e2e6035e712a77c14caf68dd0b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "67784c42396dc8bdd1428a7701e4bc7b182affefb7df04301de381ae76de951e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "2ab8d8b21594a808b113d65c5ca07569261384ef5e6d2ed322ccd582bd0f25b2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "82f081bfc1b6b78c053c96f3175a8290802adcdd16e9140c28abf834c439168d", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "702415b541cdd50eb79a7b3f0627949a6b54e5c936d5af16fcf9877ea1c56997", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "70dbb63d918193769b9dc7da54a47f16cc864a5a2e8606e2d178b58dd658a6f2", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "f2eb40a759fc74a4ca6368e8559abd5dd16ecac79d6b5cfb82c1a496a65d8a34", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "77f625acc3c8c49be5164d6cadd8b851bc7d8f415ffb5a0715d1a828a0258cbd", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "823c462adbe9c985595f39da3a817396cbb9fd9fccf4f96918f55824c4d00ee3", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "7d5b00bf92f99d654d7e84447655b673e36b76f33784e6af471668eb14155acd", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "59ce7be8b1710944f34def1a682d723c85fa65f6ab266279504a6aeb9ed2f4c4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "7f01364a72f821c6d4a810ec323af4ca6c6a18d2e8813d0990a0e1874e1a409f", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "7780b61f3a97957b51623326ebcf9e1e854dc51c171c206c7a3cce04dee487b4", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "84b01014fd3bd9dcff753d70a95fbf37028f3cd58cc5d49fb116f12d5b1eec11", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "bf32b3bd3576cd9c1fd596887bfbacb80b0639799e8da2f62131e312393929ed", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "6fc41949616e81f744751a9c5fcd4af02e7f424ecef41e8fcd22ab2fd1151113", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "e30680fb0eba0ebb87e2667c28e266de60e1cbf231ade1b6e82a9ce27ada61ec", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "96d2a9de66ee3f6d5233d6d5c93c82126a1f3feabbb595cdf148f1884ddb1f69", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "3a971b7c3aa6308528a6d2a6cb15f1530239262abd143298a40783f1d4a2a278", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "efd286dbda3dc1550d7a514c4c2df45c780f9ceda90fd9413a98a3ffea855e5b", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "289a125865b15c8289b1a30ffb1744739f38202ddb5f9863c422db07b2cb24f4", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6940e7a0e51b23129055da4d1ebad96cc9cd589d53027c80601147ee85cfa4ac", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d2d853f297f00418f74e5f394612db11518db573541ad5924850ed3b719e099f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "99b06ae72798f6f9989f093485b71fad66f58b0797d993bbf875d27bfa9f5029", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "c1d1b263eb0daa485bbcc4ccde0f761882cb813f010a0de3220b563af13a7197", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "e3884bc0f5762c1fb2a2633244d346bc7d274f7b3c829a00fdcd662052eac162", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "38c08f0b52900425f79d71ff7861ba1d20e279a7dbdfb41aeacaed2ee1d0f8d3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "32b9163d5740396bba5919b79fb6d8988557bfc1eded39f2578179d7efa0d674", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "30b939f33d32ec17f7473bfe21a46104f456fc9d3bd770df3609b248f095d0eb", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ccbf57126671f0fb540f81ba6a4e8d6d680a71a011bba8bb8a1dd119e4af9955", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "3dcccabfd278e90c60dfa59b2c7c0cc4b17ef587b32c9edafe65e088088f4a00", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "6f63a42cf130745d44bddc9f97a94cc586e9e253eb0b439178e792ef86e1dc8f", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "0f381a52636e657c5c4112388b932db87bb683433660595d9490e8665a616159", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "38da3e394c8057dff0aa847d0963a43be8481967a63180ea7303512b75592c2b", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "9bbcdc772debfb3255d63823a2e4feed1025d705db5f839b3b8c2c2468c32379", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "c914137a3bd74f66b62c9b9683881182d0310344e039976956994d4b440d3d91", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d050ab86678779fa813a1f5308a0de55f8b99e7bca1bd25bb3b2de8fb63fbc9d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "87fa86647dd887caee69c7a83947673b6ee2b7c668c1e930ce0425f2b250722c", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "367e4769794c381d443fdbaaabbfd7cc464f57ab10feae8a56ca17b6c8965e77", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "3926709b4d63a5cca2b6229219644d9892fb49e0deff01196f9a13d6cf86cdb7", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "6da3f680d9b93e4b9725a4513982fdc6fdd0753e44fe88a79cb8ff68183ed784", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "0f825a045691b22de3241ae86725cc1a974269328b91f80d5ff48f6203488b34", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "9abdf7c2a6c85739572da28459b8710b2cb995a129259d56be0aba7784de0d81", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "24104120f12577affbd0f66cf0c3d2e99794d4ff7265662f9dc45ccb53e03272", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "219222bcea72e27524a965bfdcee1b335e6b02c83abbafe65a82581278797d15", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "b5bad911db44163bd089b97aaaf89be50e25132fa983520c57cf41132a75113f", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "c48d4bda52456cf85706e21b8839f5d6397a4d8e08fd85530e9735786bfc0a5e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "cbdacfa601b58a863604353716e4a028bb8104be3056f2adcf72da79cab92c81", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "f42e6724ce238b6d34798f348be634fad94b81e35cd1f3ef5a84ed5ff02049ce", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "507cce70cefa6028e49308ba625ca24949a45fedf1340c6d09fa639578bcb3f9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9bd42dee6cfbd37ac478c39da15e23f15ea672b6a89625e3a71b43112a3fb156", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "910f7510b952de2941125637b3a085ae39350de1eafbf3626512922058e6235f", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "9bad80a8ec1f8b25ece568b045962924b6d229a3adf73ecc5486ccf3f09465ee", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "cf720414f58b5d693b919de23d0008c1292ab2ae598bf9a83b9ae46965d8e5e0", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "3959155c38719576807fbf5e246ed423211f83419916e1488f00eb3b54718341", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "e38c0520161e18eee2c44b8589c2a69a38981802781fe4c8073bfc7c75fbbfa0", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d5f774369c53536c364b9c70157c53a220496af8cc6f8b3cf6fe7bbe7705a083", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "fd025e807d6edba082e83af765204be05638241b2ec120ed62743b9b907a70df", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "6025c9d3c20261b77a8f73dffd34184565e863662086e06cce7998c3fce04ed9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "8874934246a793b86a492f10d433b26134e122c3ccbf7b2df7c262af19bd29e1", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "a345d3748930c898cd9b0180a9b83cf0238ed2655cf478428be505b824284abd", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "38b5fb9dd01c30c98bbb3276585f7183f4c3ddcc3b99b568ae41dcf728e81269", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8a87fc2de8bf516e7e94f0dc767097b51499a5d28d4bce519da08f73a950c196", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "16ccd51925defbefda4ff929b091f5060968dc7fc4a2e52290db886e5355cbca", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d6c3bbed65b53cfd8de302eb0754243b676d7b8fb438798960f66353d1a3328e", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "e6a7ef923164880d43a108a82293fc1069121fe6a20dc5bbd39b7ae8b90dbe81", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "6045b6a70787fbe86a4cdf45d9976f3c21f59b9074304e5202828990d528b617", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "964d3279edefabd4c5366040a92e1ac819b83c391a72fafba231f083ed198348", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "9f5090d8b7f936c4d82b722b5eb7759062d64fe65ce793b0ee0384515fe30045", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "0ac0d659040f413c35ea3c779fdf508b2a7162c51b073a3d24cd9daeb8c0da2d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "90d55745117ce65b158f0f7bd2717dc1283f5c3831b6388152fffe989c6d1918", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "35cd4ae1c42e9dcb5db81708ffa9308159ae4128513d83aa6ce6bc581383026d", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "68475d8742a107426797987f614793247e6503674add80dccc0dd10a2621d33c", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "60890239d190cd14fdd449f07f781315aac8f554d0ed4416d2875eb13f5a0cc0", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "cdce7b5369f1f1f33e0211bb531ac5854dfe201c5538c1a3c9698991d1d31b8d", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "1b07c36761a90d18883830b5247c54c410ce159ec424b7687507dbab0b7af6a0", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b5f73138161fd103583fd82b41694b48bcee2739f4a848c6158693116bddaf22", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f5cf22bf60d12e9f7021f821465353c38e407c550b3ca331c54fe80037b58d7b", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "60157d1673073140f9c4739ba3f2b96ff41c73fb21e9f763e05906ab17fbb88a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "c1d56f785018cfdd5ef591dcdfc2b4409a21f3b0521be34696c96434032afa66", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "882db7e7f1c49e74fc0d39b7b688332ba0acf4b17b36ba69fc153739df0c3ba3", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "83c083726ab481da8b71dde820a9c56f60ce7df4e1bc1d072e9be9079a7b9855", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8fac151e55e87f52ad773a78a436f2e7f5dda174764239d3b45b24733d05a33b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "fc3a345dbc36f02fb219f5ff7ed7de9294ce78b563fed92d1d580a74eb3a394c", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "3f9ba5b0a7cf263d7c27e14e577c1717eaffefc311467d2ba74d5e82802304d3", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "ef378c9e51cceb637d8070ce8996f7d1bd07c4092a99360dfea021cf854c1f62", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "a5378ed9c34faccecb354405814d9079d459a369c4d55908720453f29d72c61f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d36e5708a30cd4500848b41462f3290d3a83520dad475dd62394fdb5a4cc188c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "29afc7cbe24f5fdd8f8686fac44fded9b2f70df7075fc9f7103d1c762ca8feac", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "f27e399cdd6ca5c75c94c404b12a522dd230daabc701c2e20999c60497cbba96", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "862bc92e79df16bf5bafb2a5d29e579817495d768eb74aa0ab700bc5d5f30716", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "6d59b53a1f938a61da55634c7c23daf33bd578a57067692c9978932bb66d82d9", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "daa97614ca5151f695bd9606b003e37b670209a3f0ddea43aaa415e36e4ad6d3", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "c3cc872a81608a2966be6dd6239811c9adaa8ae3ba1a903abdaabebe2b02c6aa", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "7880fb27633e1716219e88d88f020769daf47dc1d7806c6b8b27403eeaea9341", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d28c9d0b5c602fcea5b6879602ce7071778bdceeb9cbdb4ba44e948fe8628a0a", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "84f863bd397924bb9b982cd105a451dbda5dd1180b122402fce3b7edcf807e26", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "7e30ecfcd71ea8f025ac8f23a7ccd15ae9ee07c6d8987a2cdf0b927e9d5977e5", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "078553b2ab6cf30afcb81f0e91b07822bffdecc20afebf29eec536b3c64d1701", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5b78fe227f3c53724f1bb35120e96ec357391f26433fb851a50047ae43e48aee", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "a904508e534bf7551fb9db58c18ab5a57d68dc876454cb13d888a882db4ca87b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "ef8b4e8a4b2e70e1e916e344520f6db30e6eff79349d0fe77b93ab2ebdd247a7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b0551b8c0bf8f667df64db3cf05427cbb9160493c08b12ae407bd60e9bb5cc65", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "92a0a0bb4e6443e21b47e494f8f7ee0f229c9b899112705c0ce5ce92610f827f", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "488b7fd87dd97501fc642dceadc074a6dcc0b206b018cc32c36facd7a38915fc", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "2f78df700c4784fd41501f7dd7695d701a7e5e725aba06c17200f0468c22ed79", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "58324211ad5f88554b5fdca9685143b2ae697bb5ff55fb9ed2b2cd91094d1561", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "9cf511506264450c0321cad0fc0f84554d51b7fb7d343d07575a36f23d691726", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "6ff723e74e82b584488d507494f24806abe3875df0ccc4877979d22c078dc755", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "8401da717f786fbad4be4109605ef3ef1046810e23a224d01212dc545bb02ca7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "88cb96f6bc22e99c6b811284e0b92074a76bc0bde8149c744f88d07664ae331c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "e2ab73f581e11f615d32a7c6aaaac2093fc987da863022179b7b4746dcb79bb9", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "9ea3ad9d2f289c529d584fdefa7c66240fa6a633eb90110991fb7e417cfc086f", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "414648d1d29e0cced41a57019db3419c6d727b6c2352c43ce7da9f6eb2a8e38c", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "6408f7f7b9c87817e8d50f6d0e0116b74e0cb4edd1870c15b089468299990378", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "dd51528864240386e5e6a90997e61a092f36794c90ab4945313d67679c226bfa", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "dafeea17b919cef7385b3a151d4e7346137ed99c6066f8fbc03c42e0f307205b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "ffa830945aa68aea1920878591e37d72a91606b5e23808b102777d3552534b41", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "861723f53a53839b74382b783f6e68d219af692e90196f13aae11edfd1c1cc01", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "2dab11fe331e721d5802da07bbc81110385e59acc696a38b1a4a3b55d4f6cc54", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "0e988100f6e1413e944fa63f688b6d26c536cfb518b52b2a0fa534177778779d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "d50687722a17d80b39a2cf4c84702903f34e4f24e16ee0e05b27b231baafc0a7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "321edd1c3d42cf7b93975c28cb1ca55af276eba4a1312a25f8cf0faafc79a8c2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "7c47c74e99c4d9e4bf63b12487de5915faef38cfcd3970767c5b23c5c4f18d92", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b50e3d9dcff88e414933f85a0fac4cfc0f032cd2d1b30fa1a102fab26f22b444", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "04d5d68105aeb4234b351510d18e8ff9b692b8aac3c5b24e94d598990cc78ca2", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "325e0f3886c8a9eb03cb22579c3eca206f54698337f0c6b2656b6493507c3464", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "9c2db450488b51b179c05debb605ae0ea28c8e419b374a8c167bfa139fbc064a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "44002b77af0255b5680cf2f9d4244ea79efb4b520b1b1855cd4bec58823afeb5", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "19c4f813d3e31c9d0d6e16681262533b1fb96d006e6cc7323a0318af9705fdcf", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "dcfeb2222cf88de1a3011153c3bf40efbb1fbe453c57c1ca1ffcd5facffa0854", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "9d9ead1d27b237f4ca757936bb617d7178e1e84d085c80628dca86625ad8c324", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "ea2081a70f2c4c013cd0b135599864d8f5a82707175abf52873c33957588f292", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "93f5483566ee44f13f69d5e4d58bb8e89e39c7aa9b7e4348277f7244875273af", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "16fc6d36e41a4379151fb95336f5bf94d40ffe78d577af491d1e3dcc90959043", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "e21031d5ce72b6a8a782ad43aa2eaeb3fcf897eb806285ceb75957673e7835ef", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "82a440d6d640c0f9e6d76ab9946dda1913fd77a34541876637d9c9d783e4c9a7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "3ebb808359c3a889514b6495b490d7bd22c8ac1c52f9d39b82b48f2f73430c7c", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "51a3f832e710e385efc89ba681e8bcd331361177cc94d8db6f5540041d7690c8", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "774ea4eb45db981b2b4e490be1990b065ea335c022cf41e8ae3a0e63fb80dbba", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e1959b38ba91d7a7229b0bd0043bf22d87a20ade7ad61060c521b385a9b7c5bf", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2f06bb4719efac0fce23160e985003d1ab1a43371b9fb1926c89340c9ae8d0eb", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "0223a2984bb944c9cf06c64da6329e9cb630877f8f05493fb5bfa5af1728166e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "a2182c11c5fd2356d363a43ab0644fbfc27a01de6252c2e222fdecdd1e7b4a6b", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "bac858891a4b383a72bee4eee7366d3d71f23312f2c7bbc47332cacfd6b46c82", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "0581086b235a092cf38c750445c6f4f1da45f823b0bb6cb42cdbe9adb1ee493e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "5618839c74f1ea0a7af5b13b8606052611c84e3865daa1060699767a3475ce6e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "c17b6d1cfd3dc10a8c59a25ba429aa0169eb04d96e78f99f9eab3c222d0b5693", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "5ea94fd2cc9f78440708260dbdf71200c407c991a20462d986ce054453d08ed0", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "9ccaca663c3922ff90f75f67732009012664c177a12be703016b47bbec297ee1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "cf78128e454f496f5ea16073d294407c43725387f31f8e018dd03dcbd9f8c598", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "53dc257ba335c92e40cea1b15c82643e9be7d4cb0c9f37785bbd2d6992b41625", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "dc0cbc3d9e70185f18271757ee038ec06a1d164c1478398169dbdca427fbaaf4", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "f890e1d0beab9c9bf1e7ad06bf727f88dfa12ecf0db6eae0c050c2700fda8645", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "707031974b691a74e50152c1df049e35b4bac090c788f0f9c9e0fb10f56c7566", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "55a6d24cd3e5240354d24e8b3ea30d11c673f22d24b5b3f3fbed27f7224bb73d", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "80e21ab8e9f14966a4f743c9b81be058e557b5ee16c3845637d68c1331856b26", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "2c7da2acc3a11c7a34665cdba4f2f46574f8926d7f6db5389586f4283f11ebfe", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "0ce356ced5180d39137dae314510a86ee5ee86401d7d15671af390bb1df9d44f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "af047cc068ff98f914a1a91a88a7955dc80941e2ac9eb40d0ccb1bd844d8a028", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c0cff535c1f6921795b4ac6372caba872d9257aa7367b3aba4aed12cd4abf45a", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "acc88379e1db298e9a582dd8c0bb502e98f519a34a2a1220924cd18a7d6bbdb1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "96e738bfa6a05de182bce332f938b77df89865dc0097807eea385af005bfceef", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "083d4f335a94d0ee73170b4d697a2a183bd01c1f9ef538f54da5571a93a1faec", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "883097fae4c2e7dbcda253b646c6b040696fcf215643024f75173e819611e491", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "9a4ab45203bc0f0505c208ac44e6cf364cf31014e93170a8b68c1eed6341f55c", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "31c659f936a4bf09232dc0c1c3d5e41f2c0655c64e42bdb7b1317c983e833bc9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "638ae7fe3e3c1736e9db21a4b211777a4ac2df60360e8c81936574074a4ff0e2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "1c23e1d3e183b6f952f3c3ff00b327280dd06d0a7ed497f831732f8f867ec906", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "2f36f13dd3032cee4b441a50b04a468a29f3ca828a890db72c6dd47fe56321b4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "858a9530757b50e59d9e1b1221ace08ee5dee407f495944e493122c6ca30eca9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "76d4700b226e3e625346e453c37db30c77851ee3c803ca13f13ccdd176613d86", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "7c03da3781301f6c2e3929409b89eadda2db59f888d19d730af5c1cf1ddc77ed", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "c2507d8f2ea1c374a3e91ae4ab49f8b6211ef1fbe3b73d338eeb786cdf1f26dc", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "a5a3c2a385b0795003699bc2ab22202bffd72aa6f74b2d0655c972ea5a4d1d32", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "6c5116a9a37fb5369ab058cfee9ca4f71bb39dc3e92200c448dba9bf543c046d", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "16a28fe670a3d3c57f87f333ab42757e40868555ce85fb935af6f34f6f010778", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "1c4da4c01dcb1bac46ceec777010af32b209f376160a4c4560991c5ad3e33535", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "0276c103dfdf646ab97b8386cefdadd0529cd9231dff0eaead057e03ead3c61c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "838649689f9337acf112595accb6f571c42da73465e30c2a6b639fc0256897c1", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "808c3989eb3ccc52868734831b06921c0c35d503ca35a0684a05579eb9927f96", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "d7f73e111b19a2cc1d80f2a3f719ffc26b390edbecf3066d2ad2b4bd04659850", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "084c00090bbab049634c1def1a43234b5d30730e1908ab6b5cd790701f167083", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "2e8d0a1c4379abb3bb09978682f031fe4a97092fc3565d9fc4771dde87aa20b5", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "5daf5d2ec6c78ceabc9f3b8d325df17e9c611b57b68f835617650619c5e9e9ef", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "808b12edf06d6f3a0036fcc61f8f482cd0eae0e2e3a7bd498176ead165658125", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "f3700d2123a069b562af02b6596b7b770292a4c36fb349fa1a6425672955d174", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "fbda6db25402a4d1c43d84b880ad56e2cc167e71b4ce47cf8c149cd5f920f85d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "f29def4d23c755d2c5166ff1fd5088879a029d790efc8b5217cf98614d55e7d6", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "3f39e71085e113c7b882473cd7573059fd4ac3a0b326251be6c12108ed402c3f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "c2661ab3961df22134231cc509c29062f35e3e820e7902192609d3340dc42a05", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "e63934ad18673c90ff7c707b7215d8261bab2038463c7a347a1daa4dd207b740", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "d053ca41366e90085f9ee434cd85805bc11f03638a2000636311ecc41141c8f5", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "3671b8ccbb9a6e2041be0ce7d836174e8f73c40375e5f2c6657d9c38dd20decc", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "8718fa225eb0312caa0bf5bea099c2ca2dda99838fb5981f1d5b44a817f6abee", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "a31116d11a9a69b772034a06cf590bb9e968bf575487b2cc01c3ad1ef3f4162c", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "5fc58cde4920c6df2e9820c11570160da1a30151bb1d22704fdbe6b9b606a71e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "ef2382ebd74324ae2e82460c16e240da1759976b2e31769c849e9fd7de0936a7", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "6a7558aa5221ce9099d18b2deefee1e29bf41e7e5c56bdfe4f50f9bed3967892", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "6854b165ae2e2ea4d00becff80282c5f998395f0392a66abeb395e48367b9d45", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "2d72b484bcf3088d8c193392de149bf2a059ff5b39e82cb0744760a0565a42ab", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a53ce8ac791541ec4b2d0a229a9c7cf324d13c52875f8be5eabd63a74deed322", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b231e3cefca5c96fe61c8d8612bcba62762d440ae5b8e10678cc5464fcb401aa", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d59cfb271592d1595259074505730774cf6d5a535a29a28a72a6d2f451f86a47", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "14b01de1b5da8e88fc7903c54afe7a881f18f2705a571e375c81dff7f67e9a7c", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "5172c7e0e1efe6761659a21e8a1a3339deb6ef40f6004e3544b2d1a9994d2ffc", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "54eaa9a5f5867a67e54f395d5b3897cc20e37eb9a3080b23b3b41d09b7572316", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "a3405af9f30db28782e0c7acca4357d757e8141cc8e64feac7b42edd0e831c5c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "baef0193159bcbf64194ffa09d316861c1e01b7df44994b6b03e9ba2d2d445b5", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "2710b4d81f1c09ba68e51d59f8b7a62522fa5d2f6652be403a81e4a4a3eb5aa2", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "5650efe13ff3cf580c2cc1bc53593ae493df3c8a4231113e859647bd88343098", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "016805bbbbcc88502c26fec3da1665f70f6a772306a8cec7ca5436dd6db5ae2a", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "e0f2a845b40381db7d584a777c1b5f03f84939070bd855f0ce50286d7ffdc387", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "a4725ba5f6ffb63486882eef7478ee6fb74bb0c1ccea96dc286cecb18238ce8d", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "0054ef1db001f1b7400d64264b3cd724323a1ebe318c5ba3236be2c07df66fc6", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "98ca73d58e4610cce3bc270e52f9a948d44aa2b135d806e1d479e02e002f0ed0", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8ee59f9be6d9a000be5a2d624d3d0a1af374f378a3a6bd660805a4c4ec5c358c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "7e20dd1908cca9a9b2beb55cf8eb457fef344c6e31f30d0987d95393cf816a82", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "aad9c3f0aa35b0192de9656eddbc0bf3a8c6315a8c58a98eadeb67fa84ce0b15", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "09461ec0d905f4f17e745fa620f0ad31a9e6874b3c36ea907ce8eb4fb434c4dd", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "381b5511eed0dd6fd246dadce9d995324e2402034421738dc16c3cb8b41ae465", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6a8ff5aa59a54b5f764fa56176b72368ac48702478c8c6b758c5a86a530731c0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d10086f5fa15533b3d2619bcbcd9fb19f9a3797b7419681a07a66b3e4f2a529f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "6af933b946af5ebd85c009de6d7627daaaf1e5af771b6fe17714badd8e44d668", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ba8f9e0eeb81024ad79fa57e355dd79446c4e93fb1aef18b2efb553d712d46bf", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "4a7c00928d1555115234790ac18f1aa3613e42766edc251fe4d29a2cad9b02ce", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "e8827dcf1609b982645c19b533d7a17d9a7a125258f4c5723173a1d2fc472941", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "b9ed2f38c70c45e5f88fe2d009207dcb0ed58203cd5898a533a73a7a078991e9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "32529a66e43e197a5f9d40817563045750f6e36a945d495ffc212cc76ad9ab03", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "9e55a627742dd8d9f7788e7738c945b33089b97bb0981da55c60569e2481f05b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "89bbf009c1d630afa9f0524d96b75ee965515306dbd651d81aac6b3a1fbdce12", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "cd6756b48dc4756fcb588b4300b5844cb5d65cc1d987c4bf5347c13a0545a081", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "9e12c8634a545a517bd0fc80536133e0c063494cc5c62ebedc2e1099cb11eaa4" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "59de1e326c7b6fcc5431dccd0f63e05265d6b57f91f2fe488ad6b471bd0edc2e", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 33000, + "training_contract_sha256": "27f686282981093e4bf39aa71163e8792c27b3cd753c22213bb66fed6478235b" + } +} diff --git a/models/ppo1/0/_CHECKPOINT_METADATA b/models/ppo1/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..cdf7b41b18e55242711bd127b50ad0ed30539b7b --- /dev/null +++ b/models/ppo1/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786225451065973161, "commit_timestamp_nsecs": 1786225451784439009, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo1/0/default/_METADATA b/models/ppo1/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo1/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo1/0/default/_sharding b/models/ppo1/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo1/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo1/0/default/array_metadatas/process_0 b/models/ppo1/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo1/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo1/0/default/d/648c91d28d47c7fdf4ef44a1fdd03ce7 b/models/ppo1/0/default/d/648c91d28d47c7fdf4ef44a1fdd03ce7 new file mode 100644 index 0000000000000000000000000000000000000000..0e0428e80618fc62e5954135bec1b1df41673929 --- /dev/null +++ b/models/ppo1/0/default/d/648c91d28d47c7fdf4ef44a1fdd03ce7 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c1b4b3c8e99ee6a4fcc9b2c359fb194399ca2f63c0df74cb6b698a250ce8223b +size 112804 diff --git a/models/ppo1/0/default/manifest.ocdbt b/models/ppo1/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..b4c6b8bd2d96f7d5c6194865454a0f78de45bee5 Binary files /dev/null and b/models/ppo1/0/default/manifest.ocdbt differ diff --git a/models/ppo1/0/default/ocdbt.process_0/d/0714fab6dc2ad102a2e15c141a03257b b/models/ppo1/0/default/ocdbt.process_0/d/0714fab6dc2ad102a2e15c141a03257b new file mode 100644 index 0000000000000000000000000000000000000000..4b1a203a494f8d5cd524894f631d64fffb91355b Binary files /dev/null and b/models/ppo1/0/default/ocdbt.process_0/d/0714fab6dc2ad102a2e15c141a03257b differ diff --git a/models/ppo1/0/default/ocdbt.process_0/d/a10a646ba0743686f0f1c30ee71174d4 b/models/ppo1/0/default/ocdbt.process_0/d/a10a646ba0743686f0f1c30ee71174d4 new file mode 100644 index 0000000000000000000000000000000000000000..970e98d09ac07bb679ed747d8bea6c27ac253f9f Binary files /dev/null and b/models/ppo1/0/default/ocdbt.process_0/d/a10a646ba0743686f0f1c30ee71174d4 differ diff --git a/models/ppo1/0/default/ocdbt.process_0/d/af265c2609ca8829fd3e63f04fe41305 b/models/ppo1/0/default/ocdbt.process_0/d/af265c2609ca8829fd3e63f04fe41305 new file mode 100644 index 0000000000000000000000000000000000000000..131756be5456f4c59ae3f449723cf62bf62f88bd --- /dev/null +++ b/models/ppo1/0/default/ocdbt.process_0/d/af265c2609ca8829fd3e63f04fe41305 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1a889568c15d391b40c767f95a92ee4cac3f6b6c34221cbe8324a3f2f517ed78 +size 216375296 diff --git a/models/ppo1/0/default/ocdbt.process_0/d/dc664153567cef23fe40a96f24e05ecc b/models/ppo1/0/default/ocdbt.process_0/d/dc664153567cef23fe40a96f24e05ecc new file mode 100644 index 0000000000000000000000000000000000000000..cc67c2716d3c7d5a897f0560555d49656bb4b1b3 Binary files /dev/null and b/models/ppo1/0/default/ocdbt.process_0/d/dc664153567cef23fe40a96f24e05ecc differ diff --git a/models/ppo1/0/default/ocdbt.process_0/d/eb91c04a3dabe1aac4233a3752f8ae83 b/models/ppo1/0/default/ocdbt.process_0/d/eb91c04a3dabe1aac4233a3752f8ae83 new file mode 100644 index 0000000000000000000000000000000000000000..0a54963555e202929072ed00ab6e16831618b93a --- /dev/null +++ b/models/ppo1/0/default/ocdbt.process_0/d/eb91c04a3dabe1aac4233a3752f8ae83 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:26b075c7ef67b8ecd7bbb95c99f297f2303fc0de8ed2b3f847db012f23eeec87 +size 362860544 diff --git a/models/ppo1/0/default/ocdbt.process_0/d/eeffb558247f75dea678f8cfc8d41f63 b/models/ppo1/0/default/ocdbt.process_0/d/eeffb558247f75dea678f8cfc8d41f63 new file mode 100644 index 0000000000000000000000000000000000000000..e6af2962f5caba1e866287b1d7cc9482791e9bef Binary files /dev/null and b/models/ppo1/0/default/ocdbt.process_0/d/eeffb558247f75dea678f8cfc8d41f63 differ diff --git a/models/ppo1/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo1/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..d74af1f7ecd7941567dd4ea39bb68961d1f7410f Binary files /dev/null and b/models/ppo1/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo1/checkpoint_manifest.json b/models/ppo1/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..1d5c44e1c38047581e50aa40ea12bbc58f54df3f --- /dev/null +++ b/models/ppo1/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 26000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "013c10d40e4a5f6396fe2c43448c1fe9f5528964a8600d7d3f4ffb267fcb675d", + "runtime_config_sha256": "013c10d40e4a5f6396fe2c43448c1fe9f5528964a8600d7d3f4ffb267fcb675d", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "013c10d40e4a5f6396fe2c43448c1fe9f5528964a8600d7d3f4ffb267fcb675d", + "schema_version": 1, + "training_contract_sha256": "39e95a253570dc24dc5d56fa5003d002e5e7229309103054735b067c9aa89911" + }, + "sha256": "284961c36ea9ae5bb272f7c2fcf6e2ad01242d7dab93c4d7c8847bf62a290b28" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "fcbf95eee33cb2ee1c029194a363f2bef316d0f1470c5b6a1b8021889f750d57", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "11522381e3994230ac4194dc01a219db6706067a0b50c157657e1bb426d1cba3", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "2d587c6dd196c434a560f7898c24b86e0612bf9fa3a7d367158462bb6708fe8b", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "0cfd711ce4d60d5d9beafd6314de8285d3eec0812107c8bd05eda3ed1cd4a073", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "7403fa0ab77fd79a92d5ea1034fb60fa3edca3769756176b85b4aad32f2c3441", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "4294e73c596e597d5d4b5bd3b5a4eab842f801f6a2df45e3822cf8329f63461a", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "3e1f342be4f2c35f3e3c5055ec661069379bb352ca0edfc1d7da2f5df882d662", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "a70756e8b441a0033824b8e31912359b7560402444933e80aaf0a1cc6bf2c846", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "ab0460215b671a37ca3c76786061d904cb71008fb8948ec9652bb8488f8f7581", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "fc516350325e6a4c5d6482e82f2d76d8808b6c5d63af679da005a892b17bbc8e", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "8672bd7e4d03b3d406d98b596313a4afbad73714e0ef59b3405f639c5b214b14", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "fa8a02f8dd67380a6b0fb5f87c33c73239659a0aff9241857e6e5028b99fb58a", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "a8b44617f55f1bbc03014022417c2fcfd716b41864e5345eace1e56215eae8d5", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "aae828a40023d01da0b0c3b11d12a41244361ca56cb26b18c99d28cb95d90808", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "95da9b48bc88c47cfdecb42d826006d9868619b2a919e0bf0df3b4c87eb0f186", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "f02a4408c2a306c46d2382703a92794c4949c89ab7723254c3b43280df57b913", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "bd8fbf06c69bfd4bae37739e57fe3285b92d5bd42c2a0067e8ad512d47cc9f0e", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "ef8f5eb01397646bc28286a8fe9972c45d384192466200910936362cacee43a5", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "bb52e567c3dc8e6b257254930b741052bf1c06c1370d9d52a676d4b1011b839e", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "83d3276beccf2044a6acc4a50bf0d0bd175fb3908bd3273b796670e3c1707e4f", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "6c878b27b8415580456480a532d49febe227bd07908fe4ae7fbe84f3694fe064", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "553a44b0b67b22c758ce469be84e875031d7f9cef25d0cd416bbae30c655689e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "cf9c3157a93b37f18d88ac2a94ecbee170de6584453670fe7f475d9a66ea5f28", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "a02ffa71f8a9780396e64f7a29c778a382cda50b844c19e5359dc0f1fa53f842", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "6f6478b5af335c8ba088541189863776e645b1d9e58ab6c49bd4753645d43b96", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "542474a27fcad3832acb3017062050f647f5791bfae0c99d20d1fae86214cf4e", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6e85012f54b843fce265a3c71287293bc49643f99a0eeb1c4e666461b6186885", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "11fc99e92f5c4a25145ede6bf1f47bf630fb31ef2e75aab100fef070fe3d26e9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "dd6a66f5770103e7dde95c9e66063a28b57dd1c3f5245e074c88a44f100ce0d8", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b4e41210f4b6a5e2af396b50ea16b1ea78fea14fb74f288df836b0091d678087", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "a2607f6c4c24b61fc6801d12028e93669eb325232c3837dc9969d8a07770faa2", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "74b640affc6a86d3c3b504aa617d624319a017fdf17e857f8ca6deed320fa7be", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "26ab32456c7d2f20d451d9dca91651bc1113ab3cd54aca6d48c74df30f851c42", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "fdad3c19b3fac4221274d2d213ee49109f5c95b97d4e6bc40c1499b93b6c7030", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "00d271968e2697676ecbfc3448f3d1527b7d1ffdf817418e7f68be54bcf7d65d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "c1116b0c982423c404ec5ff63bdbae22bbc6fe479d25b12ecc0dc94b448fc697", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "ed574ec9c23dd23340d611e8b8d4277170f15892f844fa6f354ec4def7c462f1", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "2815f7b603e6bd40b93abdd4dc339f27eacad10903728ccb78c9b427085d23e3", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "f9bc2b3cd42565d421d964f5723201b25dbf20b5386817e3550d2275a6f2024a", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "cc0534a46e2347e2bff29ea393d1caba04faf20c7ee5bc0741b86628857ba8ba", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "16684f14adaf8ab43b52d44ac3fb73f9ca2eea725382e2226bdb06e57623cac7", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "4296cb8c6247f64f2a3c1eda0a677b1b693f09721caaf90b7b572462f0f70a4e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "9441a04a1db2273525d74be690875231154d27600152af54eda000032f7b160a", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "cbf87cdae34b51c5122dff02d0f4bf6e79320261a76859b087cae66e862447a7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "a4a883346332db9a78fa239b0b9c5bcc466e9d82e93abbe23ea54c4d6cefe5d3", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "d277e201f1853c27c7c8c6a62cbdd47c64e9f035d98fd4da1723ffd7d4c9673a", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "46bd56e3efc62ee9500776b1fa5ac7c6f1eb1d435496a525676065b4c6974f78", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "88f302158b07de212eb5a4fa51588800da45cd0d67bf557fe8c8123f5f3b00e5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "755d5ba48a4e7ba458860f10edfcbef0a5a1bee1f610f82f764941a5b73e1973", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "61377ff4eb564a6286f8d33e80d33698de374915c6c85d6397d00c9eef308bb7", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "fce497b91708a4e259faab53d92c7be0469cc2c9afa423bfeacb94fba0658409", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "60c2655fdbb788698615fe228399549ae8ec3fe0156103b26bc24762e093d390", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ff3b43a4000877bc10624d029f7c882d7e96be666499aeeb6f49155abf9db22a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "916a46cf69752e544ba8e32831e774a60ad655a80a616cc2317136c1f15a339f", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "eff377d2d4ced31705469a33336b01ecf390422c44c16b041ade54b17de97a7e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "65d14cb7af74a5cdfc2fe4a60b47cbcd121957c891097a8cad7b67906b3f414c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "2d0c7953bb6545d199cd8afab47ce1dc2b05f73a8e23c282df66ebb745c97804", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "35ae8a57f0ef08454dc764e02a087b3aabf977f3a2769e5791ebcde7b3bf9796", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "febfffdd5ed40112e5a5f73d6aa7a682f0abdb4cd6ce1ad8dfd8ddc775af2afd", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "fb302eac4d0606eeda34d969813cf17ffa3a47351c24af3056b13537845f9ea7", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "aaa93eb3f0523f1c39e17ce52f6109a8f7149e0e887fda23af47fe56f10a2da6", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "c808287cd5417fe9ef06ccaa57b7d9e4de579d6e7cc34dfb1ba983b2508cb79a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "2e0feadb6f373a9cd4188ba163c09d36d07496edb19742442231b8a2801b8c2e", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "a0d990077d6985a4adec767047cbeca4600c9a1fa11962e9755e0fa059d8f75d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "2646017e508467a840bacf540aa40126323b20766e6bf370f179e74b0191254d", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "6ba469ec5dc513fb37a6cf894bd90ae81acb2870ea08242a61258678320773a0", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "039bc9444f78b10d3889647d0f3157316cc184be746b513a505da5ae2568326d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "5a9f054cb929cd54bc9020a7ecaafc80b01474fb737027f0d453150f6deaba5e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c05ebc858e0a5a85b546a4f967624da75d9890136d8603f804c2c5f0975d3931", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "8c22a568116636e75b7e8156d47e98180afc773062256007ab291d8a5b02f422", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "dcd72281097fdb11b2b276c009bbb585fe11110850057716c334976e4a633cc4", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "57a95ccf151e01ff446157e0e068aba1ae13f285b407600f56142373db5b5126", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "c4ab1256798761e83e504f269e1cc56d570b721b9a208ae9caf151633f183d53", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "1a9c3ea8eb91084cb297188bad5bbf7181b14618a8c369764166dcd945556d72", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "6332210a3a47c78ea02d03f509293bafaa0d7b7318436ca0c6617152c4d1fa45", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "74478150d8b59f554cf7359634914b40f5675763823bf9040f4be5068a6a8bf9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "ec76a275f0860182276f75ca881c99916d57833464be267bf5c655b181dd92fe", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "e07a5fb9db15e0c4db40f985982c4eb65138a02d4b11f1f15bb8f2951c999f6b", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "526d338dfb67f7c69cce8feee42aeab035add7b778e4adedd98ad558563b6a77", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "65c22aed4175c74447f607290f76e6c4d84b2d9efa09cb71eda7c779fa0ed138", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "eb3754dd06e8a53a3b6ec39782f5d6747f5d6ab670c6cf1d9b4671ebbf060083", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6725c131edeb0fd724b95004749df1a85e6971650840230148e8bd7313210233", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "c0449426a3d00cf775f92de35f259a8006013ef58c0bb36fd8abb5ce6fec3eeb", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "207c64ec1c9cc71b52d6d51ea0f49bea83199d9a1ca152a825091e18e4fa85c0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "aa6b3ae93dc8a2a628f94aba7c07466631eda939a2ff844cf996519256bb6c60", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f994a2758b8b99bd65f76c626d9f1d5b76b6c744b5a7e4cd707938ff51bae051", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "797582edea466ea7cd58b81ec55cb2dde055978ac4840dd2f5fd08ac8ba893b9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7e6326a1bd25ceb15a66679cab7d3040b8c640c4113bbb2ccf16ebb3d97b0b57", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "02222eac8a0f3a7e76f378f23e95ad6c9f737cc2a02f2bbd834f2a0ffca105c1", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d713945f97d11e92a30bc008a376634cf8dea3953cfad17cc252234f04f54a8c", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d5c9742f25e8a6eaa25445459e8526576b74487855fc745554430046445d6ab2", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "546e4bac1d5daab7995defcb77e2a3460a2f9471a8ed964c515ba74d034709e6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "0a780f96dbc4fdf8d93bf4fc4118bfdd2fd68420ded90904b810e6278d63b316", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "98026cadd93656c8bd35c3790f88235794a85c5beccb18e32d742b1a2374e94f", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "682e7392f087e797e8dec8e9194cd23c2baf5406510bfcdbf4dcb1e548819bcb", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "1abc75f3f4be7ce29d1196bdefa12e6411e413504247d22bce3450caa952d213", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "4c4b26d87e6ad5995226387143af2773b4298294a3548b0a118f1bc0d37b0b2a", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "541f22c4092d7a2889feec28d0a6af998fbc1fb6c2287ae0b98424a3863e51c3", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "99f53ea945e9b7406b256749a03dd53f1ee6653cb6c4868986e9c70035b6f14a", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "d3d4fd3e03b2335942c1aae7c4cbbfecfd3c83d8f71f6f8278cc0337f569abab", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "272bb24212affe3a7ca4b91acf21e2c4bb8a304f78b363bf0bca56544ca19d28", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "ea538a92e840914f454dec96928713244f8ed614fb4e0ccd24cc79152b4e8d31", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "af0676043f244a11f86b18f365244058a61af43d257b65e81aaac14046d7f344", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "889a3c80b60dd37e2ba458889cf4f48f4884cab3229e76b576c73e45f782c2a8", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "8118b105cabb0042427bafbdeebccfa54545927975f2e6fe61257cb311da8297", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c006489494c2ad6120f912224de23af83fc802e6598f74cabe0afc03325bb73a", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6dde52c95f285e6551aafc5f87bb5c76a3a0f4655e7ab15941b7f98b28048a4a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "4e42a824f61ad3ae94a61107053cb81c0ce93433c97d2e433bb8d4fcb7dd0e4c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "3bcd6edc2ef202c0ae468d29d96f3bc8875f91ddc71f76a916577e535615804e", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "ed5cdb47538f255619aade2cebf19f709935544324be4c9934f68bc5131d1faf", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "060ac8315426c44dbbf650767443e647626d927e237eb5d495164692f3b51757", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0c378c64e7b76b3102bbbe5263dd3e93424c20e3a0c95bb153ded15176be56f4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "651239a97e2c23c839ebaad2d674a3a7e51a76c1fec490db4b37929e941c02dc", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "86b2c81b8dd1eef96a3ee26d475e6e9ff05b51edc08b943c51cf9a55663b5019", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c9d7704100960cfc9ce8ac6d29102a984c0ea427894602a1b6f9220a8b7d2402", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "c270fc37a3cd2fdd789b9cdb90fa57a65481f5589c4aeadc9f3a2bf39e4edcbe", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "0f53bce76574fecac0e2247c8435883cafb7c61f3e7f7d0e8dff811896acb80d", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "1406fcd0b65e6ca8e31f8bbccc8860a52fb207116bcbec96315334b886f5e429", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "edbd5e00b0d677c1acf2102ece903c9f47cd1d2464b3871e988c10946be94adc", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "bdff00e7a9e4c850f76b61be666a885ae8228eeca10b4ad24e4760d6838ae8c4", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ac3c21d4d2382737af373b2bdaafd25d70e5ecafe2e7599f9f0bdb1ad49a82d4", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6f92ab4e8e28ec1f2fff2bd36bbc25732be5b0319de69d42c6a0ee99d0226186", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "8ee68f565f43b56d38150de0ec2beb089eda3b2dff0b19726e9f9f118eaecbbf", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d82ab62b4b42d86c2dc0494e4ebaa9a4113fff0fc65004ff1553f7f42b455689", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "76421630675b98ac63c1220c723d3a82deca4a39f2f3609a5493bd997628ae51", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "fd0a7f8cb2f897769ec84b92c81bc46612d59c5963c3c0094554a2cecc9b979d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a11cd316505fc8e3c88196231f136fcad449335ea28cf20abc91bd10890ebcab", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e1bfaf21cb95f2a5b517b05632ce5280c07f260548bc8afbf4ee7c2594db1ce8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "56262a22e0da3405474f98d35cece402eb60c52f63f31b8b61f66af732765feb", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "240b7ca37b81ad9abcc4f898dc13c31119613e4c89b3ea3f5bc485188da16fe1", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "1b6a1e232f04db82b28ffc0dd3285fbc942ca1c10d75bf2c0d15fee81d30b312", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "093cff707d2c73714e7b82615b859d18e980425202b6bc2f670b24979a9bdce4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "00da86e71d5e3a6a4698c3214b32848daf74d903dcee741d310f47fd9a0c2a00", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "7eadde43482f8997d24e9422c9b70ea30fee2b78ffc1d53276a5b1fcf33a0a0c", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c4b3b2bc327e717e9f23a34432d433bb334cab7da9852692c610fa512f9fe86f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "7574ce3a06e72c56c50b900524fb1d400cde79e9ff2d9cbbdb3818454e7a4668", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "3f33424b96fdb37d0b308e4c54058260b5c58449b8ada396a0152f70bf5cb42f", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "e14ffbaf92d41cfe3f334af0db9622a5606e3f3feca6c76b0e7f6b86a8e5527b", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "9568b16dc973c5ece2abb13733a0932953eb8429c5c8f25af53b0b8632a8e1a5", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "d341ca2d3b6c2b85d53db04ae8966de05d373de14134ed6c299502bdf6f6ecb2", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "6efb7c327350d5b869b87294d1765748cd059508f561dda7ce6bcdc164d13a72", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8dce104eb308e0c1a7a0c719f4932ae2e6fe975372f638625f1610a292077774", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "d0b50d255094d06f445b5858f7d492a02642d5f939dfcd96ac40cdd7060ee984", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "ad4fbc57573ddfeffbec6eacd4237a579643ae7e9a0b58beea2d69f6d9ca784c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "fe59f7cb3c4a010c004192c7d52273ac1d838b2b33383409edcff973485d5371", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "637309df6424ac70694a9117cc462f15a7f8fc5bec14b63c252f4677d2f3fec3", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "73bf7fc2171c198b3a9ca95019601c7a86d5be0badf507f26cf40e47292a7f37", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e8e0982b68660abe6382530a8a234773c695ed908661fa8a1cbac2bec03d97c5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "de03549d030a2337bf2725fbf3812bf09efe60a2c62d1f997b93ddfdfe9e9688", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "2523e87a91aed5ea629b8a3d81b3345c8fe076e2c7d035f4680829164b10b1c2", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "5aa58bb8e4b2a0f82822b97c8effaba5275ae33657218d44ced8e589202e573a", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "5018387f515afda096b818b11ef37e98d5f90eefd931b5e7e23873f672d1acf8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "c891193e310072ee0f46ff32e4422449b9df9a74a00256a1f3cd846b4029bc50", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "879925e94b5075343340bc99a59f0edbf1bcb6c9f812b4fcd27b8acff4a452fa", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ca3ec620afe66813543a3f94d7b8a91a5bfab8bf2635b5f06cb41c69310b6bf9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "dabf738e28833becdb2cc546829b5b41500c852d9b08b9c7419cf6c27453e587", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "70e02f5228323e3d6d4d4a9315a37d55ddd6e7eccf04ebb6d3096c7bb0901128", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "da20bd1b09aa417f5edd558c1e965183e5c00616b164850486152c3bd49bfb50", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "4253b5458717c59dcea27744da2a3e96d63fda25771f8f8c8941ef277b0d9ef6", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "a36bb1445475e6ac41763a249b2cb65b9e2706dd0dfac091591771579a661884", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d434de96c159c69b54b7b7b69f6bdf21fdcba3e6ff97eb5557bc8c6b42192a93", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "3787c1dc15eb0d0fdbe6875fb80253ee8292008bcfbee3a0b6a6b7c038f77457", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "004b07ef8a7a64d17cd6cdc2257d38db115fc1a4b3025b47a58d3fae057076d7", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "93c5e1fcc5f1b20b04f46acdfc674c9bb219a7a5b4c4f50df8cbc0ec447b6d85", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "59c52400cb52faa57362927aa0c0471e1803883defa8f4c6e2fdcf21c41d77ed", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "1c748e8c2d3bb20f65f372decc162c42bf0247972adbd7f988790f4c79092ede", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "3e0ad6664eb821c1a5a7322b082098bc2ccfced4ed22514fec065b0596307f2e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1122835b1112503abcb3ae716652d4676493b225fd24f5462cd435f5f8a6a414", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ff3a3cb9e6d682399c2f10a02a544fc2d18fe32b36942843079e32feed4d8ab2", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b8b2b5f5e905cc842e1ead061d96337467d83a23fd019caef7c8b57e73ef86e4", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "15fa853fffb6af592178685a3068abd7f9a2d79e29576203c244c12ecf401c3d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "28d2f386430ee7bee7934d092ea34a58d21329fe6ae6aac94c63a145221fe6d2", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "f26479dc86d21e9931beb778baa009ef50b9194279ddad6552c69ac9caaf21e3", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "edc4f195d7c8360bb500a76b13d5fd5fcf3735004a549a7d983eaa7b54e5e33e", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d3c4b9a34b1e59189934e4379155f4069c106cd804d6ab199dbb6eb29b1b9184", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "05cf9b86aff1698b806ae8e4cd860e7d0f511d92ef7b3dcf8df25c36599d4798", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "443d7e0eb09a16327070041fcb1766e5e0583347feb5a6283548e20958acdf2a", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "a01170ecc713fad7fca58335c2968d21f9577ece29ecc96d77baa7df599d14e3", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "63c0213d6a14e1a598878e1c07ae89d4d120240cafcea5133ee79f35ce287200", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "661a63cd20d2a3ad01995e634c727e4cdc75b9790ebac9062aa4e1def0cbd0f2", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d2446140bedce05edcbd0bb6f27dd40993fdb4547bd9f526ce3efc248ab0656f", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "06c3ebb389a007eee9cfd587acd401645d05a4ececba2b4d1f03523229d4eef6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "211c4bf91e8b056646cbc68299554e5ec87031c7eadda44f3ee2f792655ec74a", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "85b2da7f68b8b0bd5909e85e6ef7912ed4de4586bc494a8cbf9e479b873d134f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "879e55ff8d32d2efec4cff77d3193a40f214f3019cd011281a2a9152384cf672", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f5717f6d16dc06217ee6240ab48d8e8f9f72102887ec99e79a04949b03854bbb", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "277f481e5d8e92f00bf8b89c480ba8bb54ac6212e63142c9c1ced8c2f4ef0b02", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "2d06c31629c7458c9bde10df148a88f41b1a8452ceb1998fed7fbc1f20d01054", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "619b255dd882670fb84c68463046698dbcca8ea117662373e9712dbfa7d5cb7c", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "2ca5ba40261d2bffc75f8f044615fee9f31326ccd881281819cbe14dca867e16", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "4d6b9f2bf75a9f2387b17bc6877f1d11bdf7b8a2a07bd6a39f9617a9ce514424", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0c8584d9cf1b6ded879c5d683c355d5803eb11895410ef283ad0a2bcdc211b74", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "5e1955916fd21bbf7129ca7f4ce187501ea9056d6dc48e0f22b6382ec6b85918", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "ef269fcaf19adcfa4163380a79032e32611e1cf6daf2f78b17b4a4b1ff21bce6", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "bf6d35d0c80013d57fb64b05bac8ee4e916aca71a4f6c8d4d4dd644f61b971de", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "32b6bf1de1141aa9b7b934d75e19d49c90fb101039e5e6c50a71b2f39da71e90", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "4957ce3d55e64b27cc7573c39e5d4418e5707b9ae663da80e2f831a81d0af45f", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "ff6969018c8f0df5b798119ed1dc0beb93c7020c730f26832ea9a6e19edb3d15", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "77a3d3dd2cbc9ba0995db2a146fb0c7af4618e67ef781d14ca76d82f6cf92b4e", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "ed0e5ece182ddb34cf78ac9654fae88c7bf50573b3121dac039fb27eac24b697", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "21bdbe55da5a40058ca6bffb83f9de32f10b5a183f1c72ccdf6c4009e9f68bd6", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8a20dbb6595989e29a15b0f1ca4dce9c7bb0f7885a533af86b6166c4f0cb0dff", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "5d5fc697c03a7f26d2c63487b1829b51dee86d48407fff460580cfbfbe94b6d5", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "f01109ee99a17ed7d03fa43dcaf1f1b1b642f272423e06b470fafdf9095d5120", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "81bf955807f5f43ca32495cb9eba026470ebea1e9ee42f9f13f0020f3628f87c", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "317c0900de705f279732d274eeb7862c85e74feebc7f0d6bf1a602e04694744b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "20a6223160fd5eaff0d4eb3f15eef9a633fba6e79c7bfdf87d1bd63ba4ac9899", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "cec51dc260b2789c8a656faa6aeddee0492ac91b559853a6e1b1d7ffb71c100f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "6b085a7d36991fcb2f56b3a477e7efaca1fdfbeea7cca825a8c500fe8f593dac", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b37b1f4c6009bab633ae7951e631652e9c3299af61f5cd7c7ce813443baee03f", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "894a757b005d91352e59151665c8fded55732ee21c5327661844797db063d6d9", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "1a77999ea9d1345dd0b5f071070c7fb22ee18cfd9c4ab9dcce50b48d916853a1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "73b1977ec3711b00503d65842ee503f56ae6766a6ede91d81fc1528b72403e1f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "efc5410b34a758917337b26ab3594823ecfe63ff1de254d699c6011543efb95d", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "160406229ebeb91175b1bedfe6d307b77fe281c4f16ccf2a742f5d489fdf7fd4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "6aeb93ed3b631d5ebfaf4097f7a2c7d79c4c35c217549aff98c722ff2e950781", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "82b1959b0e5fb5ffba536467fbb33ebe651495888b41c6c5999a1cadb2a79d94", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "c4ca8dce59a10926ca95f8c50f703c4c74ef98626c23ca518a5916f3215cda88", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "f43f4a09ee07a4e6b4bb00b3bde723a1419cab6238639f9cd0009a51988c5c44", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "178345fd75c442f331ca06e7a63459102962a959dde2cea9131881377868071c", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "41885d5d9f9a4e30384641ed562fc9e1f79b1ab4ad154969663f756b8093efda", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "a939f4441e23138435277d57ddaad31ac6ddbf927d1630bf74e99ba54348b624", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "85d47b51014c421ccd103d6dff71b19ee4b120560edf0f7f0fe73c33d4e8a632", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "3b8ac0f56fdb2ce165a48e2bbe58a325477fbbda0ccd588b4c4b613cb98fa48a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "8b694085c8c891cbef7a980d9bc76030005bf921c87f497c05f781df852eb7b4", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c8ab6bbde54b9108dce232a1146e0e423272be969a162d60dd35113d3c89f340", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "6053e81b96e2054ad538667715f1dd2a023d63b631eb5cb7f93165e087678e45", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b201b9f64e7deda8898d46a0146af4f3045c772dd6717653c3d36d67e88ecaa5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "dff211543a884bd53c495283bba1fbb08695bf2582274c15d30405bb65180c1b", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "1b2870b9ea386c68ac3960aa69ff5a2ef310c5f72842f2779acef03ad3a4d080", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "887452a629fdefe3bcca7deb72b4f5d957b27eb821564e3949918cd4517f6916", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "97cecbf777c72b117e99cb85a7b3791a63d9efbc55aab652da3b4772953f87a2", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "85a17632cc5a7d751ad4428e5d24e4b2437a8c4032997bbaf4c724455fc3b61a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "6187e58f01cd8ba331384423ed8ac7033abc0ce4c91ae7f3d67844f191143754", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "8061b0c496b0cf5d4a9964bfdb1fb06b21a03549545441c1eab5a80acbd596ff", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "854d741a7de5673f7e0a00507dc302972ba8de6fc883bc81da5df9f312a5d5b6", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "22584274e7e832dbd5feb4b4971c4e1c6eecbf72cf701184d72c6f8744d81690", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "88fdc5d480cd5c03c0d1728ffb6a2525198c6b5574038c8098af9e130fa83015", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "eefd8c6a1a128c6467bd3dd9543653c116c727f7ab1a62dd33461f82d0ead8e1", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "89f8f1a3533bea2ec3212c1703b0ad9b81e16de0a424fe722c67cd5be05c60c8", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "a7c011b2a257c7d9592db77b2b4462067e2c2877be8d875ddf6d2ab9a6c4b3e7", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "ddfb7f930eed78c0fa9e88e146ca4531f7f93c18ba437d3c4a8572f5de3183a9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "7c2730502896d9f9884639406c4b56e3f5d70f7957bae10d633f0f947e24b97a", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d486913d04e0164755a168ccbb525cd7fae33840e008151fcaf986916a95b782", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "3ab066f1ed3f50a6eab10e9a592ec1d8cafe1b81c2e957439e1c23f2bc3910e9", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "3d87bcacfcf6524b2f325124e100b5fafae046c4e9fe4afc2040ddda4c7baf8d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "edcddc1f27c9c63c84ea7168a972495dc3f03039a6f566712c6d8afc000bec4f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "c0bb783126814826e304e789d8e579391ab77ac5f876c2deaacda1225b79044d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "96a4f779e1c1ce7ae2b3423c925f306a0556c052f187f7a3f26e2cd41f26e186", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "04d5d007a56a456945c1324de184755e09e2d55adcfd9fed2034ba40ff7d65e4", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d3d6d5e427cd7029bec8c95c79e7680cc50eb96a5adab2bfc73510ed65930d8f", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "43df0230b8796a04971457a2fc0b20432fd6861e729903eba008c4fd303fd66d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "4579d353d38ac38f779c14b988f50a7069fa87ae9c5ed12690721f52cef288ec", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b5e5a7b2a67cac7da6b0038d4389534e2e4e0bc989f317bc77fe3457abe51180", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "99715f9b0eaebde0dba34fc7435c41d31ca8f98063a45ccf33d0ecf64bafdaea", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "af29903e6aff5d92d62437684347f8b09cb0191673b61db57256fbb6f862f191", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "16f2d8f7467b812157fc3f6473a12744d95e06fafd10d5117bfd8fc9964fc187", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "5d1cc061fd9ab6a368c3b4670b1efbfd3b183c72bf828828fc147e56a2024539" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "b5bf5462efa3db7c6d6bff029a9df2fae6755f1622aa1a39fb2e8a202451b411", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 26000, + "training_contract_sha256": "39e95a253570dc24dc5d56fa5003d002e5e7229309103054735b067c9aa89911" + } +} diff --git a/models/ppo100/0/_CHECKPOINT_METADATA b/models/ppo100/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..0407cdac30b498cd87447b5794db6780acfd6b34 --- /dev/null +++ b/models/ppo100/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786225450763849264, "commit_timestamp_nsecs": 1786225451522329111, "custom_metadata": {}} \ No newline at end of file diff --git a/models/ppo100/0/default/_METADATA b/models/ppo100/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/ppo100/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/ppo100/0/default/_sharding b/models/ppo100/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/ppo100/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/ppo100/0/default/array_metadatas/process_0 b/models/ppo100/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/ppo100/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/ppo100/0/default/d/0300e5ac4b49bffe5719d9cfb4762d1a b/models/ppo100/0/default/d/0300e5ac4b49bffe5719d9cfb4762d1a new file mode 100644 index 0000000000000000000000000000000000000000..82406bd1f93f70a3e1db29f442797a2f5e8d2677 --- /dev/null +++ b/models/ppo100/0/default/d/0300e5ac4b49bffe5719d9cfb4762d1a @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7977f2ff2beab7ac7518f82248b59357e186ffc6f850fb6fa42d86cd8d558803 +size 112892 diff --git a/models/ppo100/0/default/manifest.ocdbt b/models/ppo100/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..90ea6fa6ab4554193fbf7e2b0004964e9a9726e7 Binary files /dev/null and b/models/ppo100/0/default/manifest.ocdbt differ diff --git a/models/ppo100/0/default/ocdbt.process_0/d/3c2bd48694a308b3376ff141afe1a568 b/models/ppo100/0/default/ocdbt.process_0/d/3c2bd48694a308b3376ff141afe1a568 new file mode 100644 index 0000000000000000000000000000000000000000..c1544c704668d7e5c53400a1f8e57d70abc9858e --- /dev/null +++ b/models/ppo100/0/default/ocdbt.process_0/d/3c2bd48694a308b3376ff141afe1a568 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d9e8752f34ac003dfee90a34c8d6e4430736e00361ed0516c58eeeee5a133505 +size 428994560 diff --git a/models/ppo100/0/default/ocdbt.process_0/d/417b2d123a9368637e34a59cad750a06 b/models/ppo100/0/default/ocdbt.process_0/d/417b2d123a9368637e34a59cad750a06 new file mode 100644 index 0000000000000000000000000000000000000000..970e98d09ac07bb679ed747d8bea6c27ac253f9f Binary files /dev/null and b/models/ppo100/0/default/ocdbt.process_0/d/417b2d123a9368637e34a59cad750a06 differ diff --git a/models/ppo100/0/default/ocdbt.process_0/d/8dc8e97ab9734180b7d01ab555af1a32 b/models/ppo100/0/default/ocdbt.process_0/d/8dc8e97ab9734180b7d01ab555af1a32 new file mode 100644 index 0000000000000000000000000000000000000000..4b1a203a494f8d5cd524894f631d64fffb91355b Binary files /dev/null and b/models/ppo100/0/default/ocdbt.process_0/d/8dc8e97ab9734180b7d01ab555af1a32 differ diff --git a/models/ppo100/0/default/ocdbt.process_0/d/b50d4837ed106bcce6325d46d7b977ea b/models/ppo100/0/default/ocdbt.process_0/d/b50d4837ed106bcce6325d46d7b977ea new file mode 100644 index 0000000000000000000000000000000000000000..e6af2962f5caba1e866287b1d7cc9482791e9bef Binary files /dev/null and b/models/ppo100/0/default/ocdbt.process_0/d/b50d4837ed106bcce6325d46d7b977ea differ diff --git a/models/ppo100/0/default/ocdbt.process_0/d/c08f0db736a2a786d6c2b6a5f8be193d b/models/ppo100/0/default/ocdbt.process_0/d/c08f0db736a2a786d6c2b6a5f8be193d new file mode 100644 index 0000000000000000000000000000000000000000..83e08188bfcb7f45ec362f5fd0755f185b6cdc2f --- /dev/null +++ b/models/ppo100/0/default/ocdbt.process_0/d/c08f0db736a2a786d6c2b6a5f8be193d @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f4cbecf6daf77edfaf04070d82012ae964161961732ccb62b898e89b40814383 +size 150392832 diff --git a/models/ppo100/0/default/ocdbt.process_0/manifest.ocdbt b/models/ppo100/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..51a8a2ab36aa6275525c6c82d82a81df9e275270 Binary files /dev/null and b/models/ppo100/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/ppo100/checkpoint_manifest.json b/models/ppo100/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..031386d894c77b33567e68901a54881ec516f7c3 --- /dev/null +++ b/models/ppo100/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 44000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "b6f3bba96e16db1aa9c3e7a31ff990a976477c800d6ab92a39c8999b4b370770", + "runtime_config_sha256": "b6f3bba96e16db1aa9c3e7a31ff990a976477c800d6ab92a39c8999b4b370770", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "b6f3bba96e16db1aa9c3e7a31ff990a976477c800d6ab92a39c8999b4b370770", + "schema_version": 1, + "training_contract_sha256": "43848407b7ced83973d7571f0ca785d04c5c8a8747a442e2d1a42a97d0c78cfc" + }, + "sha256": "1c965ef6c1cd422d744c9332d365fcaacfadc78994f3be007d371871bbf434f7" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "f772b92c8ccf6da2d3771a60b3ac7c839253ce9b2f2e6bda02a9b83a0de2988a", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "78db6320d5dedceb9d93761229555467efaa35b3a36265900f8967d9dc5d29a2", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "3ebca9eb31af8a8701ca5c3efc8d7881cf9dfdb199eec630de6a992d0e67d49a", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "2fe275d2b9a5228eb19df8257172987f8436ffec706f353bfef29cf810f581db", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "ae37ed0fdb60146d1dbce200fc7435857135665c917a5145691a371d1cf14211", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "c879c7919dd2746c11d4984009e3132ae4e86f1429dbd07ff3c29a6b0a5a6989", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "56c2f90f1a14889b7d6bbcc8b6b492ff9a2db1d8a3c6e7acea882efb7c2e851c", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "6183a9e370cd337f5708ea2c90dfa5af835ca7a9665938652ae2040f31eb2216", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "a8661ebe1b44d4cbb71d325cfe7534bbef71926d4c7bf14c8811d92a0426373e", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "430db36a8ab3b9b39639a822f5849f33f18b54166ddddcb3d902836f5ccb0522", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "76108da081dc97b19f876689014561ea82428222cf0e650d6624b76052c5ad25", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "5a894e9b6357f8f5b8f2aebf3b3e7c4d79d2d02b31423a1d334e069e395ea9c6", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "4dc4213b35d501965f5c8ced2432c0c0ad6875181282ac00896d3faf76e6fc88", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "4d52ead9d41f0e3bde394152db60fa65c6a3f4df31af92f9239e2e98be837bbf", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "0a69a737ecdded0ce24a0cc84a7ee22c38bb8776e67a7b3a2e0a028bce7760b5", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "6b1231cf3272e06cbd18a11379e7e2ca623077f3389f4d879b1469b2701d6199", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "58a103cd5a9bd09496441033e49a06b929df7cdeb4f681db6091af144abd232e", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "822f7421f1c1b8ac7a7bb97b876f45f8269f05d7bbfc6192a6a5a9e091605d63", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "d6664d3e34522acee41dac92a37218057887f61b8d39b56cc546a86fa137c4f4", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "74d859f0b1d5d00531abf02fd51b81db031f6b3de2b9ba8026435a867a3efef2", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "75ad837a0386258189b152e2e852542cc8869a974a8b0313bce4c8c953794683", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "aa60d7ebd7717d556019fe5e54b7ec6a02ec62e8f82824bc7431c383aa94a989", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "42ff1560fdf444c45b49619c2c67d128ee8072466d3c937a05dea72aed1cbf74", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c2ac73dfb30d805874a5747dc0eb58695b2263282c5877114abdf235548c4160", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "9ec25e364b2a3f854c9a11d21778953b265b73c1880ac6ca89bace6dbf2c8d3a", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "78455da29ccd079aca2cde9aa8bf9e352f2170103cef784c7c92fe4e4113dfd1", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "02c68c1f37de8f6dcd1356b24fae1554869073b996d7ebbd0fab971e7e86f6b4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "23cac7f7eb9e832e2845a8c4d94ecf46046b20153544f5e10634a5a01960a29a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "bd27d96d125f96b88efa82dcb099c3493cab22727edde7f5c114522a692d5a70", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "424fcc2cc21176aac5ff5ddb5c7045d1a444aea2af261af11936868100714a03", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "ca80900a3438e909730d49a2d7b6b80123405ff326eb855a2e2498dd7abe54d7", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "5ae914fdf5d29e6144fbe19e73442cfc4f110e88b813c7c7940a9dbda05a946f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "655207474611322f7e774acca6549f9d290d0f878cf2d3baad51561a8ea15346", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "70347faf0b8edd30c912c9f50f1fe4ad806fe0102afa0bd641041b72bbb6e8da", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "545633d9e41ed1aae218b13b2d954028f98e3fb33e23248b9858f51624366a1a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "383f4168f963455997683aacca4c4050080fc16da4630d1eaaaeb484eba34f9d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "93e3fc181676b9936a54c489ccd641f90cf7d83ccdaa3d6bc4116ed3596b842e", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "c39a73e236fcecc203da029c2a6faddfbd2776820842c7a93bf26a8452579d63", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "414bfb0c8ac82196ff466866f96ff51eb27824bd72ab144470015b6adaf2610e", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "ddb53d6590332e15e239e708521f80d1c96977391d79abcf390ed417e74382fc", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "189e1b90205790ea185d769a017d7bfd1ec553f5f5b5dae8b2474fd2e20e7b41", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "bd69ac99124ba91e1401f74c9659475a560be17e6053a79437606398a1715322", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "161f95f4894817790f41ecfc9f21559f9fbf13f5923ffeefc1f3dceca6d641d1", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "39e7d465df044e4654a44fd7710748b6c1f809386d39c400248f288401138880", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "65908101537be28a3705f7ba93aced67943d1e5b635ca0dc73d353f5b89e33ad", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "4b7aad6b1ec1426502be0343149f4a02ec5267761eaab815e1d3f58edaed8e40", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "44deab0d952698190c9d840b0bf34b963bd58406fd7d804567b9c2aa8628ba40", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "84d545621ccc2bd482a71a02864588f5f8faf12fbdd6b5a1c9648f878acf64ba", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "21047e51ee9676cf3179c776daeb062eb83c55a66784e29adc6a2d672ea7647b", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e90083d3b5a2d9b396eab650d192beae815495af864f37bb4a63701426875851", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "905f7be0e0f843d7978220a219e524e964c29a36ae1872989ce255cdd32e6f58", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "f3a775ee65085b17b65ffda2b8cc89d43d87c508b5be51ba6e2a9832a7322806", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "7e4a79c97c1d2ac6a90d6a7f2e05ba52595d87959dfb85e3db28cc849d759eb8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "f98effc9f51558b9b3832b3ac77c7573f9c68778a8308d1a2eec9e185cac47db", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "344e546a0743947b269622cdb6d3abfeb60116c86f1dc90a7afc6885ffccb356", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "769786959d450f3f947d345e2692e91694e8e89d649208f71386a6bc69515149", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "0b970ac6d8fbe3eb7b1f738a0be2b3d233357427554dc375496d011138894300", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "4e07ba623514a29c8a0837c1204c51fc55b66a7b42037a5e10a5a4e4875a759a", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "f20f17ba04d39ec787849941d10249e6a24ef8ed3387ef99dd12a679bd0f4e1a", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "caed893d7c2bf98ed28d702f8af22eb70f80ab2f852ccf50771ce5d315e85353", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "525e019dd4701829713d30e1ac7ab487f5ac73aa81ca9f0f263d5edb1b05b4df", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6fc9a951da00c5b1909dce2b356af56622c8a4b456b2da317e8887b561c0eccc", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "12a99b11780f3f8a8056c7da78ec0f011fd58bdbfffb01f93850383e62e3601a", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "84c8cb7f0660f3254196d75f982a88cb12cd84a8300c912b23b09c93ad7d82d9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "9904b627ef126d329ce08fe8c6e7d77e7b020629ed1044720c461e3a04ebda11", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "15f684f0f1800d3c4c9a774f5bef99166606477b25eeee886f41c0c5b19f8839", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "9dd7aa843e441fdad54d7eb45425f49d6f4f1ad60c90db9e2d06411a3cee0161", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "f2622c220ca4df3135cd078ed6e5f9521b41a538749b2e1e547b7a8e4641c882", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ca8996ebec2b09d004fe054e1c333543ac475221a248f116df3affa5e553648f", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e7f55f334369a230ebb854ecac67159010819a66ae6942146d1c9d5a5338f8e9", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "8ae5e62bb75db9d779140b6f3af214e629c49de989a2d67af0c22f0f04e60dba", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "243378826f834efd7de5ac56d3842224f3e978639bce88bf5d68981911e27b96", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "8e30c1fc087a5c5af1823bd5586e32b201c0d4014f29058cceaf5bf00a264e0e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "5e1bf508125ab2a2a90886bbc169f5ad4491d737061f9824254650f1ed5ea2e4", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d8b9ace34d0116987efcc7252dbd25a04549a6459aaccb032d1d60172cea860f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "8ab1ed804380a01c68096b6516937dcb10f3bc2a479bcfa5a9623aeea61f6fcf", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "0d5f39108ed5d2253c3dd989ba1e9fec25d5e89a4a71fd237a4449029fd6efcb", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "533a1c60b06b20df5b0341885945e1f46feafe88c711e03ce22c72f38e26fa58", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "f306bc1061a2756c2382fed803151acf2f18a029cf2efc9135aea63e457a0cc1", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "4797c93661c99c516ab00552effaeda3e52a1e84bb4823021304ee084a077dbc", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "02aaf37fb44da4ad18a07dbf2461c742466be2996cc3757efbea2f4a658ff38f", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "e3571a715bf2efb4eb07329cf136072ffd7c6412c6a0e2650ab3ef7318cfc56f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "9a55bbe58b25de33a8ef744575e628568ad9d1f59e283be79d0b513a64decacb", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c0d241a0d9d11784c8b715ea88ea5db504b3ce4ac6bf60f8a884b5db32ff0b52", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "d879ac45817195038c2c7ad36754cf10439009c24dfabe7063e13cbf1bbe5ee8", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "ab46a33b6963b03ebeed2a4364f02e467cdb05ce2564d80de0680c3f95248956", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "4b4cccbf50875e44019fbfeec1d4c42057ac56e60b09596b2af525ce93bbde0c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1310dade46fa17bb430d40370e386af93491a08e02d8a4b9db5b2af1d88e2d6f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "47f215c61a279d04c713a95133204af52138b453890ede4a5c001dbf67f9d99b", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "db538aad5f0d3e57c5210ea3afa1de5bc428bf29c79cb15b5526ed22bc5128ff", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "03d201b5128f3cd3f460dd3b9e0999e2208331d537090550f31ba5949a1bfdaf", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "56d0ba720584ee7dd8aaa9349d4832ecd25d20d8392afc6faf88922828b5f57d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "3ec85e166d0bada9785435ebbce89272c3621713eed99d637c43342c503eca12", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "46df8d41df3f5f4df5126bdb1beb329fe6d103272c19074ffe8a9cee6fad30b8", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "395a16ccc30a4ac0a27c857b1269ed4db949c768470631b44865b11806ce71a5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "4fbabf0b24c80927bf589c5e835945a6ba40b209af7c49582a35f4b59042dd8b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "a5e8b446fd435931f5089de578d6302ae6e0f5682244e3bbfb0b2f8ade41e37b", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "31e7d3600f2e339b1af5782b5867125da3b16eecd63866969a05a5900a7fa7bc", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "e6f2e843b1761734d17e4d3190aa19ad2aebf63e4d5dfe71b7b92108d4c09cab", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "27898786fb80f5f8b798070a57747b5d9d7d759112040d1b4ad9fe5a7da445aa", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "0c37776a3b610bbbc49bf7b3c554b53d4810aae337f9627043e0c214584ce3e2", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f528054bc68b7db024248060f07828ae95bc0e93f1b97a638114f7812cb50c67", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f52ea8a03f319bff71cbf61db1d6268b79302962365a9c886ffeae91d694e3b7", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "6cebf4d31a1e0b9fcaafcece221595b1a3004ebae237cbdfe36fbefc1dc5a11a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "11e2c3e063cc5971e6c582eb0f12998ebbbc8108aa620f2ab81f1130b1b02fa7", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "089cad067cc0eab872f15054721653e3de634a217f05c25fcb9acb5f4815808e", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "01aafb3d7adde7d1aeb4b86d9e034f07bd9033b63715aaa6d049653582be9153", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b6b29bdc188ba3392c087c952c5f13a73b0c51503106383c936a51450bc3328b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d5dd1fadaa13bacb0de14fd86cd39b4f2312638f923c33b67cf4c26dd295d10e", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "33a6576cfcf72df881daa51596faec7ed40297a2e45344ef098483dd7e49dc7f", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "84069bab7eba64aa9462c3fd9ae6423ddcc664037b016924b8549464bc73d98c", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "75af7794197878f6ba8f33b74bbe872b8ef911289ba7ee86bb1b6373b47ee4bb", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "7c677da385e89b161eda6fa0d2a9f2d25389e3959fd90b2d9ee99b495a5b7beb", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "a540c37b8776fd15cf302e28bd2011e311a245dc09f4685ca1eb2fb14ab7ecea", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "0f067a0d80bd9c72b6cfd354a23cec621bde80c0dc5e3927ea9cd5905913053d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "d7338505b441b02cde7b752d8ce4358819e97397b0c084c3300d09600ea58698", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "17e889c1c7b81b4abb09d2e243669105d001254487dfdf5297e61fd73efbabf9", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "c8d953a57a6f13939dc22f9ca26c8ae50d950daad58444e733a1bc42e809e371", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "fedfba6dcf4309ac8736e6455509cf2b121f4a0a6c4943b2ad859153e3e7f146", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "731e6c160617052e96c84ae109c737cc0ba122986ebd41f9bb078a573a7ea9de", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "c40e56efd30896bcfe2857474bba0cee0bbb77c169e1b66f48e755cf779b8d2f", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "56b8acfecf30f4813b46f810689016f37d82f0a8a2cefa5b1f413ffb84da5939", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "dfd339dc5a764946906e5c9a766d2cc9512d42edb27c43e3c4ea73cf39e03f8e", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "9da93b25bd2d93c2173a0ee06713a7ae6ddb7cac5e42e77e1495cc3ba172518b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "ceef6a0ab747299ab14002eef5fd63b8e183fd51972824e520f8724364b9c068", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "f1a4b27c664329f81802612353384d48ee868099d1d57ac694d1d78b758d0070", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "56c68297f8d096bc191881827ee34c093c6bf96931a9e896c56b312a1a904ece", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "0299ff1473e3b8caba302e435da8d3fc150da454f05f88fd35e22fc11a42e23b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "50918a72ee90c877e654d74820233f408fa573dc8cec80a41c53b3e77af38fe8", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e7804a011c3178a067a404fc56016c2027da15ce85215c1d732053b2cd869e93", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "ec1fb6c61c9218fdd830657ea20ae182f3629216ba60fbb55c041023a13c1256", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "95f394a3846bb682030915b039422674c0d601114709252620746d7a550a7b0a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "7b1c811b334fc5cc2c9e06de717ad5cb6107b25e7ff0038a611c6107447c9ffc", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "7a657959c711aaf5ef434b7a99f798c87a54e488fe8487dff953dec5137981c0", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "be20281e9304683a02afd2af45db61b3e01bfcb36450b7290cddcffbd7c5bb83", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "5385b8796cc30dbf456183cd91d81e1b8b6e2aa387a8adc86e5207d0b816496d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "87575b91cc4bfffa4827309b5956974681d1fb04ec00995973293145ea1dbab1", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "61456e9ff050eb7056c49f2ec07bee88d2cb31acd5358dcea086b2c8b34b07b6", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "652df840eab7df62af6faa321b71d11b0c3a74c5f2e706921dd21bc124d694c5", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "c1ba363424d22434235c3d4af46b2d1ca09cb3a20ed09b1418a9d0f6695546e4", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "9ede110021135daca511a2a0d9d0c8bfa428dcb369557005839b8c4f4ca8a880", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "631d9a355a8db81ed37f7a2935799310333190065ab9755348e9e6f25a78a6e4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "296b974197ed938fe917e2536d1280a5b68d9ea7adf1e721c895ee35607293e5", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "b53f0fb9f935c4a2a3b94bc74ea595695f298d7b11ca5e17f03745b35a890769", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "456afbe5f1caa745f02662b9863cfab392e617b99f356cdd7aff117d9d94b078", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "bd085c37eefdbf4907649fca537e58ccab240c0f05ecb05d26028d6064b7e007", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "cf43d16cf58f323dd0d69803d9b42ea4eb21a669bb8a1ec3952bc3e0d16f5ac7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8d4a4de1f43e7f18fcb8f25e66b03ac93c4cb95bb4c1c7e1661ea5ee75b5b782", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "5d4b5aff14049a32b30d35f04dabdb65d5c544c366de3b855d1666be2ecfe19c", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b8b4a85b1183540f63fd38caf1a9adfd65d9d9d449993748058813e1eb54f860", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "8f0b0a2628621b30dbec9a0d843682851a29e84a5fb225ca4ce1dd0599435800", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "361eaacad19b466e3ea2b3c601b7a8b078e20ce52372ad60a07b1a4c107c6bbe", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "a29839614482f3ff68d8efa75f55679dfa78eda47bed2a28ba6e0e30bffd6e11", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "323e7fc7a3b542dc1b0639bd70c92f08fc483faa0f0d450f95c17b10dd6a68a0", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "6ac6892243cfb1daa93a13bfff8735833c07f2b95d803b21ecc57db65d2879ee", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9c3425f59f37118537f55d24eb579ca4f5a275b5377fe12804ecd92df8a0f827", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "d5dd675456e532f76ae7f127193313ac374c4859067481188b469fce3db02b8e", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "f10fd10f86f6e2e393f8c5516135107350c907d4af6087c4908cb4c040571c7c", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "f40248055b0fb6fbe91e78e2100fb21fc084b1945512f86139d9cb45db89b918", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "39ca37a7e60b0df67b81d1bd323fe1580bf1483ac4cd3543d2ce9c36982c5413", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "697f3627e2cafdf8645bc3762b1bcf0f9cff4fd5644941b9edb10a64b26b9825", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8d1086491af76a83e10aa540f2a163470cd3a17e5cd966c9b54f501f75341495", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "a658354bf5495edd63ea7332bd650a6110f0ea4c22b1e26904f5bbd4ca4d515f", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "acbce86c8b93d147188b9cf88f7a7a1107d820e723fb6cb2462ca9e1903d7207", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "b117c81e9ccffda23cf6cdea037b744ecb8767fcac7f10a704fc8b1303445f4c", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "2a9d3ff914718c05a266d818b1eb29b34c527ebaf65216a63131a1e8c3343e99", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "fe3fb68d3d5e6887e84641d365d92581b0b2161adbe91baac1575601c17e3d8a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "61d97703c2619a48a8aa9e91f12d4c373138dbfef9f7fc8c87b77b74b6b09317", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "62d479a72716a3d48dcbf52ce84e0893481e3885ca91e31560d7b3ae315809eb", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d2519271d6c25afa590438011798dc51755b82875bf7d555459191b6df4776a2", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "3fba301b6ba0a49634d66e90ed2c77bb9f5e80227cd2572319bc6eb8240c8f12", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "1c683d02805b7f1148035f1ab652c3e4f513ed9f9bcca1c1e53533c269408283", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d1bde1c528bcfad4891486150750c10b14deb2e55f056b6ac40d7f8de52e8742", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "59e7c56bff9dcb99cfe8738e95a041be5c3e6590e1b04ff7558d9dfe8e332cf6", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "4ea342672b61890579e52b366cd372d2f8f7859a0125ad2f7ffd241ca033cc51", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9a21c5f25e2c6d577ba1d5412701c44116ce99a7bf691a849f9a82c55ecaf98e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "9445d2cbf4ee769986a3e1388e4ab5fe93f173a3a3deed2edfce9b91980dd989", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "cda75ed3b12970c232a7ae75a03b2ae9b3bd9aefbd6be87e580c457b022117ab", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "0fd3dcd84d4669f547d5b06fa191db6220ab84212c79e4a17f82b3bd4280a0ec", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "2ad7e15c7581a474c1f2f67276895fa5ec26be166eb408fb5ac9b7d56918f069", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "571196ee00e517c6f458ffafa3f18d6935b4af41b6f95666da3fb79357970cf7", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "738fae6ff05c2807f73aa24ce237d5a19029f3a848c6a9d1912f1f405966e663", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "1754a142e4ae9b2b2ea2824e3e6cad0a3db75173d6d4087defc9eaf1bbb95b75", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "34b6c987210d68dab41d94eaf41cf5d50db81808539ad19b9a5e375d2c862c0b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "59b101183a4282274422a443e4d4e9b15acbe18cd82820573a9af5a720bdf1fe", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "23a5d8a000767530991b125c18e33a794d07d361abb5a898ab49c078d744e555", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "f2457951a901cc6b9a79214e9bcecfd826b64d4cf3d40a7df0c56b1bb513e006", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "3226d3660cbe6a113f72d4e639176a8c9f3908ea680435e27b0d2c3a0b3e4d57", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "6ae07e951a7eb7ca5bed9c7b34d92a4896222f44bb512583b7635bf5cac50d44", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "9de0089291930155738dfec2bb10bece751ec83a144cecb3cc97b857bf8e211e", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "c89bebb219dabdd6b4a513aa07809f7bed2d45ef4260cbe1535957f4e4dc0bfe", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "36f906bf1410ad939373197135b6dbf70565384600ad311c4910f6deed5c3eda", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "83fa80ff858a48cc75529ff19ba1229ce1557642e42efc2c59233f2fb123ecaa", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "91a9d02674a268cc426b545a17ee0978981d802f64cfd910823bc32595d705ff", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "5e299cdc658a4965ee4c6a4f52e3bf668e720a488ca2418f0235518438daabcf", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "d19d7d12f19a3daad811be0f70126770402e41550e040a5b60334c9b755b6ca1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "685a6dcca97f306771509f43d979a805b05138031ba43c58a467dbb58d24d397", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "e044dafe584227bb039a89b95de2fd067103b6f634eed1f1a4f1a769bc65c9c8", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "00fcf3bf8b89788900abc358a21bc7c51717bba327ce27ab478bcaa7139b2248", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "763a63b631e221ee857102b8ed70ec864a6ade3c7487ffaf52adf39b4b5713a9", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "76c3fa304725c7898b160b94484fb83aa8dd466aec88a9bad2955c0004729aff", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "8c9a6a94393c89d86bc1a4c90da45e81e460a95721be63b6f03c065665377c86", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "6fb171bd02b81723660da2b0f7693aac5b1d61c95cdcd6f4cabd6235110af1b3", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "64ca056b78c1ea42f6e4969f5d10326406e0d8ee873d8fcdc058608252fcc588", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "35e7cf95e67aaec41bf8a648f04668158a057fe44a2344cccab3c1c1a10f0b25", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "9bbbe6cc9b9f1404a1c3ae468ab122bda650aae47b5ad8ee3b11c3c0338b3db9", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "2657d8c12f6b158ebd6cd7c735363e487402637a44949307ceab112992d87ac3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "3e8305ede61f91a932c65ed4242991140f595cd6d4fff9ec6c126104b88969f1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "e78a6ea8b92d27ca80cbd232abf7fc3bcd2a3cea7629e9e17bbe0273171f4230", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "a9437f5c4d8ab98c5486b6bc0a37679b90f8a516e8277805d1594fc187a8026b", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "da5b516a53a2e7290ff5683cd76057f2b8fbe2d7df5c18c58ae994c9f7e0d70c", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "6994168c4b42a9d074a6dc5d7311dafa7c80c4d7e498a5d1c4463c8fb3cf514f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "70a68e83662cadae5392b98080a0f605267c1b4bb390709486d7e37948ab92e7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "3235d043bb11709ca535cdfa00f766403cccccda97ae626a86c745cd33d0eb5b", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "bcfd35657dbb760db6942acc775ac41cede7d460dfeb01cead696f79387b7c2f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "55fcf20414450df83cdc594120df41470d44619d1bd6ac0dd77827e165b7fec5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "1d70e1bef533d6af00343ea6ca13ae5f91328e57ac33cc0eb2b712ca0c199c91", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "d398e1bb03117e42cb169017a180a94692bf57d1801f84cd1298f979678c7ae3", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "866b3ae6fad57fb409ea67a91ee889fe32a041c3d9b7685d873ea230f8cf6aa5", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "d6e156848d2a620cb6fa99340421b87c3f948c0ab4abbb019cd012c62874b4a8", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "6d29fba8b879c7c08250c27aa327fcbd376f42ac765ab5643efa43e901e64de2", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "e74fee2a7c9d44fe865c0229eb01890b6828dd9f5e479fd363d8f36c4c5f7421", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "a0e82b344363313362e3822e0adff65808eea45b8b200618f95db6312290eaa2", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "55cf944c5c4de194acf56156f0dc5f9e521d5d9b23ece4306e6bd87e00d818ea", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "376b5741a671ca946bd27ea1bab6860ae8f4070b6644dfbac1f5d7dcfcef2417", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "a45bb163f644a41cff52cb08a2ae0c92a4e6f3167c606cbe4d54420798413c6b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "05ae3f6565bc10f463379b380aefe23cc06e33e59a735103a763f04d8b027d3e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "059b8bb977c70e09ee787f72dc014a007055504f55baa014fa546315b2520fc6", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "e036aea97fa846fc9f13080c6060abcbf27d8ff724c621f976bfd5cbafa200cf", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d345af79b554b74129830e8f15b791ae7ea74e430911d2f87cd6f4ea09e1e83f", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "ad0ae97cf716659a0c02391b9ebda0e47e6b7cf5492cf975b43e639c1544a803", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "300133895bdf9d266185c4ff5886f8fbbf1cb9b21a2d9d20641aefa707f16130", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "1f2ef306c8f968ba67f30d54493cbf1b0ee0398d00f033b9911c11484f61b0b5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "1f7056ec62cacd4d9a1a8b98704423c07f2fb7ffe88888d136d0298397eb0f81", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "f410cfc278088bab6913333cddc2ad0f4cb305e08216b940fc26f62c0c8c09bf", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "480f737f0cdaf3792ed879b463b296075df5056d18afa71d7ef80766cac6d1d9", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "3c87f46dbf2f1d9df0335ea19c3bc93abca79f6d913c57b5cd7df2be6bc8cf7d", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "b6e9ce55004d74549ad78b612c4d332554b0523de7407f55de763fc1bad090d1", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "f119ca5e7fcba8d004b80337d22cb5d5411812218d23fa902eb0b88c7c2fd362", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "c4fa51f376753680e02076c8d304cc8f0984693bf87eb61753836bffbdc21832", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "b57804cf647ce9306d1d1a055101f6e83e8d2b2d66ff506f80fe38a940f803d7", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f7967d67bd0fe7e200c17065d3fa042d96b00a1a261bdba1aad5ba12ece23d09", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "537cc6d9a22057a28b6105f92483ec8cab4fca90fd91b32e6fa89827225e9e11", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "60229b4eab4df37e8597185cf63ddd12be0bcd7467b4797c22629230cd3ffc04", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "7380eac6fc231dfa6afcbd0480e1095503f0275374a5b99851821745f976715f", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "829d27b68eb6447e056087754312c0b6d46e6c4c1096a8db225ff07663d9bea9", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "21a82b919b9e497ff6a6be71a79573bb714d080404ab8b3219e7c76eeca8baf8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "850d5a5467d5e623d4ca4700f38bed0fcaf7e9b9331efe95c4420a5c2bffab5b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "91a1baaa9eb483ee2c6cebf7885212fe08b807174c83e61c2d64f596361e5f1f", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "d7bf06fbe551a0f30a04b3dab75ab1e2b981888f49302d9fb9d12ddfa9158b9e", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "12c2c2a4c68efb05f635e1bedafcf3f4c383954e5535da7022a992b5fae56124", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "a2292569a5ab27755c182362e3b2952d394dc55e9a756f7d3af2c27c3b8fc507", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "7e9a4ca81457afb3b6c7ab05e66063dfe4a2a6a5779ad5e2f3916cc88cd59e11", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "9d72eef8094b5f76ea7fac6fe74a987cb247e344f82df9219a18d753089d1afc", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ef286fd1ad601554848e795e0d28a10fe418a01b521b1130aaacc7ea62b31ce1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "70761aff110e6914b3a783ef00ed3c253c891157d766f708c567ed2c9a19fb3a", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "61da19dca36dd3f3385a991e3a3f68c8a144ba50f665502cfae7dd7c86b084b4", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "f840a648a0208577c15247f1119a548580814ce702bf0c1f9bfa41923cf684cb" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "d58aa358b2e5ca076ff4ff32e2c2dbda0560cd9d51f1066020121b68aa25e847", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 44000, + "training_contract_sha256": "43848407b7ced83973d7571f0ca785d04c5c8a8747a442e2d1a42a97d0c78cfc" + } +} diff --git a/models/recovery1-transferred-aim500/0/_CHECKPOINT_METADATA b/models/recovery1-transferred-aim500/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..2ab127ade02d005bb9c17200e7435557201fca06 --- /dev/null +++ b/models/recovery1-transferred-aim500/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786293416380551705, "commit_timestamp_nsecs": 1786293417054796481, "custom_metadata": {}} \ No newline at end of file diff --git a/models/recovery1-transferred-aim500/0/default/_METADATA b/models/recovery1-transferred-aim500/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/recovery1-transferred-aim500/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/recovery1-transferred-aim500/0/default/_sharding b/models/recovery1-transferred-aim500/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/recovery1-transferred-aim500/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/recovery1-transferred-aim500/0/default/array_metadatas/process_0 b/models/recovery1-transferred-aim500/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/recovery1-transferred-aim500/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/recovery1-transferred-aim500/0/default/d/fc2640fe3dc160b15796bb4e93702707 b/models/recovery1-transferred-aim500/0/default/d/fc2640fe3dc160b15796bb4e93702707 new file mode 100644 index 0000000000000000000000000000000000000000..6987393615b2cdddaf210b98260c2684fa7d42d8 --- /dev/null +++ b/models/recovery1-transferred-aim500/0/default/d/fc2640fe3dc160b15796bb4e93702707 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:50c61d1ec7f080c8f58caf9aea5ea09c286abefca47ac44d49212f753d620a1e +size 112894 diff --git a/models/recovery1-transferred-aim500/0/default/manifest.ocdbt b/models/recovery1-transferred-aim500/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..61371f6c87c8a6075d319805dfb84a9822b21309 Binary files /dev/null and b/models/recovery1-transferred-aim500/0/default/manifest.ocdbt differ diff --git a/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/262859f5412e92128ec9d77f68d5849e b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/262859f5412e92128ec9d77f68d5849e new file mode 100644 index 0000000000000000000000000000000000000000..5589a31c367ee8bd07cf00471049270f644d58ff --- /dev/null +++ b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/262859f5412e92128ec9d77f68d5849e @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5f79dc226cd0a885d34007b5ae814e43a64f14ccfa688e2b46bb32999c0fff26 +size 116887552 diff --git a/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/362947cf928ae683427fe0c00bfe6855 b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/362947cf928ae683427fe0c00bfe6855 new file mode 100644 index 0000000000000000000000000000000000000000..e9c738f8819c35eef01de950412350a91c47eaeb --- /dev/null +++ b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/362947cf928ae683427fe0c00bfe6855 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6c10e570deef510408653dec66f5575b2f2faafb9e7ff5e63f2351fbf9e31c2f +size 462483456 diff --git a/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/506a3289e41cfc477618955ee791c8b0 b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/506a3289e41cfc477618955ee791c8b0 new file mode 100644 index 0000000000000000000000000000000000000000..4b1a203a494f8d5cd524894f631d64fffb91355b Binary files /dev/null and b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/506a3289e41cfc477618955ee791c8b0 differ diff --git a/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/57b7895daa6a39948024d1eadfaa6e87 b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/57b7895daa6a39948024d1eadfaa6e87 new file mode 100644 index 0000000000000000000000000000000000000000..e6af2962f5caba1e866287b1d7cc9482791e9bef Binary files /dev/null and b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/57b7895daa6a39948024d1eadfaa6e87 differ diff --git a/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/9a0dde52525302e7d178310d631e3cc7 b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/9a0dde52525302e7d178310d631e3cc7 new file mode 100644 index 0000000000000000000000000000000000000000..beec0bf9ebe44fc2e300c4f692e30b79b4176703 Binary files /dev/null and b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/9a0dde52525302e7d178310d631e3cc7 differ diff --git a/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/b48ef8617c1cec65c1ef28d3a634acd3 b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/b48ef8617c1cec65c1ef28d3a634acd3 new file mode 100644 index 0000000000000000000000000000000000000000..970e98d09ac07bb679ed747d8bea6c27ac253f9f Binary files /dev/null and b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/d/b48ef8617c1cec65c1ef28d3a634acd3 differ diff --git a/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/manifest.ocdbt b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..8d46a0b3881d6010f61c8752129a55a4597c1855 Binary files /dev/null and b/models/recovery1-transferred-aim500/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/recovery1-transferred-aim500/checkpoint_manifest.json b/models/recovery1-transferred-aim500/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..e0401559258f013a944eff0fcfa6b67a278b6929 --- /dev/null +++ b/models/recovery1-transferred-aim500/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 42000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "dedf3dcb1d0f8127dfedce6a6276e7063102888b2c0d4d9c4162cb00538080b3", + "runtime_config_sha256": "dedf3dcb1d0f8127dfedce6a6276e7063102888b2c0d4d9c4162cb00538080b3", + "source_commit": "102cd770d0e9bfa6b1fb24813f6623b84e8f38ca", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "dedf3dcb1d0f8127dfedce6a6276e7063102888b2c0d4d9c4162cb00538080b3", + "schema_version": 1, + "training_contract_sha256": "5cf2baf317839f07eefd2ed65b785e5dabf16971c1df797972e4a2ab07558168" + }, + "sha256": "74d8f079243d09fea16d86867485b6d3356982f3ae7182ae5106bd3a195d5d25" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "5998d8e611b893c3b8918ac7c2d1fbe6e801a116e94d8d564d4982f815c36838", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "488fd119ac2ef88b805e815ce679f616a17fa462ef46d414d71ec1419864773a", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "712750a5302cd735de5b44420e30522556b32d8602cc5006da03071d2d9691d0", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "b55e0a88746a559f27229335b45fa823ba8f227bf3026834e3013506edde64ce", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "46c1411d7a60eae9fcbfb14dca9ad2bbe4da0b9101cdbdab5ae90865327815e7", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "c4a53c38aade6c42e5f95dc25d6ec0b9dc13b448ff63434bdd7b00afaf5fae99", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "b993f44daf8f83d840be8cb6eb1cb773531a797a86390e575a9a179d9b2b5296", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "6ad545ef79a76d58d7357cee79ae74a7bf86e4370b9f10ae76c42deff9dd9527", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "f656fe263a931f6eff56c542a068f9fb6e96c021f871d6f7fdd6b24a706e376a", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "4cdc4b22464c21d04e929eabb2d0ed893844a3697c3a5c08f41d29105ca57fb9", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "3a7ddb92b21b266c43dcc4ff84e0445c20f22945bcc2631afa6105930dbd1eec", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "fd58c109465feab74d48d15916c08b644a62368df1e40acb6af67a1f0ccbeb87", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "cc3419c777174f9e59ce040115e0da8fc3234ada636c9399bb5ab606ee90e30b", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "8f2e88c8280653bde081c0b5a60779cf2025cf4acbac8a5c88942fa69d17a793", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "bc8942bee8fbd5e13d86d08ab3b4ec7a46afe5edc4e2ac4d421026ea819dc43d", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "63aaca3c314ff99145391d250e4559d53adaddf5eea1909bd76f99d5e411707d", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "c9badc3fb0dc537bd2ac384d792d4b32584b543ef4ddde93ffc7c3312f188080", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "f85133cd6d977e2f8088c50e40eb5edb1eb23a2d17205bca07070e8a848ea9c8", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "ac983a9c773f1f5ce66b048163b4e849b0a74ff83e489eab145a3d7c48dd2392", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "f062332f7e12111bdb11f8a1dc54edbfd7aa4483e882ad86d699493476c4566d", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "e07c12ac17f9ef031e89e54d2bdd53b71bb968ab6b907bea69bad775845cdbb8", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "6a2c24a0d7e2639b7f73abffa17834a4473327a4d1b6ce5a98538e1ba866805c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "afd1f596efca367e5bf948e43eee0144370104c2003f7f55bbfd13056c837e0b", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "5fa3a44cde1b06a7f9c5603aa4a2e97a6c447f6cc94924955208fc5e0d8e22e8", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "da571df187ea2fa681537a6eed1c5f64e13f0d1449b43a57d0db38c622bcc3c2", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "de0cee1c8fcea4ef0c2de731cc095601799c70dcb8b22871c519247773a3cedd", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "1e8111ffd477bbc33528519918458a8950cc0cc2a65654feb3933f3bc23646d7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7123c8375420eb2caaec974227f94736e037c3e8ea1ca599a00d01b8ed7fb521", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c90b4c10960f52331777122ed1b3ee844adf5f8edd9fb255f670d569cc95f7aa", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "763b6d6f8b3049135948837801413e8714fdc1226b1a3a442211551204881a11", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "bb210eed3a29e1f543ca034a220ef243f62ce05041fa8a946e1e9e35a61fd592", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "562ce0d29cbebbfdb235e64da558b78dd2ca1381c0dec9e42ab66b4d24c33155", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "4facda2bd57faeec61f63794ca4c1681c2b1138d252e5256fb9bc44681fade94", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "10cca18f5a91f743445abb8fbbcb73704437a4870ca97b4da3b36ec16cc005b4", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c52d95c62c7b17e4062d456be6cd12c0d1d7d18893c46bf22d01cecbfbfaf6ff", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "fc03a0c2e7986107417f70306131b50ac7972b670a164af20f336d2b6672674c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "d8c33e61e9f31e9032c08134fd2f253222ee55580ef77c04bae9ce3e529a4181", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "7b854270a689821932ed5e5b97b9d7cfd4deab7ad0d2ff022c5c84dfc0e56931", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "f88c53156a68548d0b3cd540d6369085602d989bd344f126506cb2f475b582f8", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "3f227b239057c1141911c01fa4d6b72bdb50e62d95e839c842d6a2428e9ed108", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "af2794c2dafa91a3a1711cd46aa08a5551d0f996e39d7fe43132d5ad3495d363", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "d0d800a4cd90f3414e3bbff2cc011cc22e6b5d56fdfb8aa383d804eb9b833d01", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "0aeb139147488355066bc791bba9316156ff248766c226c5b8d2ad6055b22796", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "8ce7c6c41770092f020b5d414ca9e4a4675f2f03258e167a43fc18e3201771dd", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "844994a462cecdf9009195f54098aa4761205632d9a11e91621431ffbace71b4", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e405e78dc4a8ff154e998c8ab49a44e8766bf2865590ec5ba1d66eeb82101001", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "f82198101ae44875ab17eaa8e0e36b10a02b2e01a0f538f4fb293600c4c4e052", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1f85cfc9c1e1ca4af705be40cccb0b0f688acb5828d1ebda09cec8655696d995", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "b4a6dd068a5d7443798f22a5aa33bfcd7e8c64d2b5bb8f9c3ffba2ace7b63fc2", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "6e47d6839b2b3c62a7866fbff1853bd5ec9e877adc60b745df0938d6ff87332b", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "558e982f8b1d98f750b7571661bd2462712fc71129bb8b3867356d3f7d905f4e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "998dbbfc079e4f8b3a9ede2065f7028cb1028d1ab2486ad4082098b1bd9e1e67", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "6dedb120befe2f20bf2a41988f25b8d7287a23bc6329efb5d9a7cb2f659cb0d2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "a3886f4671543ffbb5fe329e684888fff656e652699a60c1e44db25bc190684f", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "237ef67e46d4225a8f9cdfe4f150c68f9cfe288e7c9e8db5c3e6b7885493f6a6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "f24c45808058989bb19f7f105215ec7f267ec3b862e5425c3fa50a4e457a1969", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "c0fc0ddf65229e5309029ecb0fde8eca4e9d3cf3046f5be3b101081f171ee773", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "e5e1e001a8e18f9eb7cf4c32ca799c9828b9af84fd2e33a21f1df96acdfa19fe", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "bd7f59c451c49533c648ad14f5426dc86f49f394b31a1967960c0cca47991175", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "00debf0ab3fed86a712a135e2f17aaec7ec4e95b77c8df8cfcaad918d3da5764", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "03bd9d83c557ee0fa84b99540b20e2a3825f800d3ae691323b76e51c9bfb64a8", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "fd80ee10adbec8becdc77e3546d86445f09860ac6eb6468433541e9482d34952", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "8b691d0538bfece58c519b71d121486df5fc214370c5b5a3adc527f924239b8a", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "755220b40c1e9c51f710958c8571e7847118eff4e7b09db76d2ef7f70dad891e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "75391a89a3b41ccbb11a97e7a8a135b222e7231feb0f129de483e800e9f485e2", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "679392779ebfa8099dc404ca3cf856a6db991325151937539d87359fff4f29cf", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "ad6b401833eb6e665a6125ff9ec2be522db61753527957a9bec0cd32aaec426a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "361a198e394dda080302f2ae977a340ab0e11af853228dbb9c7fa238fd1e2f20", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "19715d15003d52be67188f986cb6b770f0ddcbfd45203faa029abac74bbe0b0a", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "3712c3f8714a9b4239c4b407fab49430fd4ae6bfe7fe5607aa5df489f4c6406c", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "63c7a73cc94270c95ddb9347f8dcdcdfce9d9e972b62027c8cf4cd77df4cc8d7", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "07de8058cb44d8d95e067f2d92f6cebf3846e743af9b20ad0129d8dea6e9bc27", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "6615e3d03dc8291b2fe5057773011176e4c74bb536c182b9a289c1364057d67c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "3d04566a35940581582548adc1918b6184b412742951e348ed95fb4ab10173a6", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "77bc38561e9b0092a1f467aa74e8dcca8f8c391528290603520e9cb70f9bb6f9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "110fff22cf247c08a0152837f852a5b4df041d100b521459259b33eb85bca3ca", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "b63f3d61643a64d6945a99672492de05efe79c1087f08bc62b429954064ab0ae", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "3c3128083d92f9f7079c503aad5f6af6041a65531ca05ae7154a91efcf89aafe", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "17c7f208d4fb812e1effec42e50591d2acd28031beb86d6ab222bedd814b2fda", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "8ff309e8920e635cc6054b9009f688eb17072dfe7f2537f3defffe14dd254f3d", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "c9d05fc4d3ebefb3866fd649c32625042e8e841d7f395e7fd0c7d6ecabc36083", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "04b1de336f20bc531f30e79fa9f7e7acfbc56e882b22ff5172a6229d1d4f5c19", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "7620e65163bb6dedce8334e30ac99f56149b16cc7ed2b7d5229874456119a574", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "5456a2480a5cebb15cd3fab5bf648779fb765f155438b55c988993191d4c3c6d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "bce3bd8409343ae71564d145384c4f422f6d5d008d8ab5037e8463d78b871ca8", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "5a0f1481eecb8bd57d7a80acfc30125cb975d0578091df8490982aa80e747001", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "24a0d52c216f9aa8d28422e99ca6080dcae42a77dffaa007327b05e9808a868b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "78b8b69967725f888a409c0a67183065d5e98afaa584cbf9594b423bc9e9afe8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "aee18e2bc53f28b99909cc31f1ac0bf5d474d7ba18f1eefe9ae4dbd693074ee9", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "0d0d142ce039d239b575a08e0d808aaaebe7e086b1ce53548b56f02e9196ba77", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "5a6373ae43d143921bb85955f4c45b84089755357f608164e165fbf0b0482cb4", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "3150070a90e2e93ba03e4693c090979b96b923c520b334becd55d8165a345ab1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "86e0302454001b800dbb20df15cba8c09765bd518606b975fb71a49b29834368", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "74d96c06d1836aa6b3fb79e8f9e71f407be372096b51c3acd2513c86d500a2da", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "e16bbddca0124e5100616fa2c1cace690d13a06c8fa1677f84e3322f1305ce0e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "d701439a64fb53c0b30c940ab96127c81af7a9b35401ac9c5e39b237ec6cf58d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "8a6e76e097600a5b3978ba406c7dd1154ecb28910ed66168eebaae346aef41ad", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "fc1d6a528bfa658976699244eb10ed6b5729df2e72de44102543924968a6c654", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "821962356ab0ac5106fcd31eecb82f24cec517f2122f6396efef1dda3a9dcdca", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "d20e79d8775553eada1cc9a7ed05d63d1be874ae6874ec0505e0dc3439dd963f", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f716ea94fb08e7d3b84c8693275cc81d264c30fd691837550228fdf3b091ffdd", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "bc4b5e323a850a8b27e3743e1d4fd43d1ed56d69a6b0a954c147c7703cb2fbd6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "04b3574afef3da066a9783f317ec456c63d7903587c3b7aaf087789e64cd645a", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "4c74f10f170b3702debb640802c56f13c335422e83488c2cc0a427368390f781", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "1f6555b2d5dfb603c92b678cabdbdfd5c3a919098dc721199fa6ee0023427e48", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "2e5966d3ae5cb714fc88c9e4e4848463a01d77190c8e7fbc7d40851cc3c22b4e", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "f963a94ab0f12769721fc5541f1ddb05c30ee9af168121e8ef4bb202316f76cd", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "b16567dfa743e10c3cf51f193a4e17048a57d1e50dc9d59e5dd6409e96d486d7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "086b36ffe2cc8ed99e1115d3dd81c5cfe28f975c02230f10c74a614ff726ca98", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "dc08dde7a12bb14310d32d5f2a0c4b2ebbcc0c8085aa3f7001dd01e3c760ab24", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "f096459efdfdb59c9ab93df5c2cb5beba38a4afad4ec6eafb93fc91af5b47042", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7cfe6886eba04fa41c8ad5278708eb8cf9183ebe439fc93670c9845a3c976508", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "4fe54bb69fee226e4fd2bf39c044b0ace1f655b9c6b322ca4a410c13c54f2837", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "6de77b765f341c386a64d1fb42a47217fc0180fd7f604c2bc048932d6359c333", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "64593a8d635ca3a33670d5159b6e3a692bb642b88de5eb1e9aac87f42e7ea528", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "6f6607ecc5bd2a3ce1ee082453900096b500cd16c1cbca720511fdc84cd8780b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "012b37852301307aa85db0f0f7435d2c7428a8175e06d8e1c363fc1ee6ea3c64", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "5dc1b1ce6e4758d68087812277f02e62a55f6b94b3af1dd1823c06672e1ae8f7", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "0780be1ac9d833ea4e83cf938c2c16247d925af7ce7080f6cd585f948916dee5", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "8b55e6ac3b81c15c68d63965b93356ac8db4fd93e83646762d2a57d140857cb8", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "2053bdf04c0b71d6f8d4a1971a03fcf16847c5bcee81b97a42e36b8f76c27989", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "bccc1f6e380925a1859394ccadb7a8d02cb4c0cb202fd3fbaadeaabbc0a943a3", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "682dd46931007e9f21f8166fd7b96333f457dd968e2ce6be3732191a6f07eff3", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "27e68ee7dc09c7c90867066c14e9ec1ce5b53f6c6508c6dec82c3885a89d3831", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "763d3eac93a7f1c9d11dd0ec7649a3b894124a26dac5f0ce3b9511fa0d30d248", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "2b046972c7a6e8a706c30b122b3b0d176d4c85b2026a401af74ca3bb0f4290be", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "89b4a55b5a91c968d459edb71e6fe0ecd2d3c6b1165ded5f2011cb800aac3e40", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "16f02f2a4bac13e6707d9621e6b26556548b8c46cbf4e36b8668e250793d1d4b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "f08c3532b77cde498b0a905934e525f706ef30cc4125c6cd375f3d8f7bdbafbf", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "527d6e083e5202c78690c0579889e6495b38b73841841b017764a90080af5b86", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "f22dc020e6441281f9e35bfd142785cc7a1b22dfb996363c042269c9edf79744", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7a9a1db41e23b3147107fd59e3e582b5d3ba7e805a4ba18681dcc80330bd73a3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "0d4040442a8c2288cff3333e2f88f9779dc0408d3d65cb84241c24cf4434984a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "29eb952f9907ebfcc5df2a0ab0c1aa71d63ace267ab7b66de61c0ee6c282b6d3", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ea0b7608c242f511723522ab2a9621dc6070b9587f65106ea0016b4b5ee6660e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9ce9828d0d4d9c3a8eabddd86d7f335526cd152340bf0ef3d95b874a31bf3b43", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "4ca5aa199e479c45e14ad87607d3abbb8c929458522c6ef2ab4f3a54d9fd9b86", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "3993b656b6acb8905a6f879b2f809769ba19bce63de318cd2125f3e37efdcde1", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "40032166ddca3365cff21c5f26ba09a9840772b045ba32510dd0594ce73885db", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "21e9316530f7a6f6976df0565d5b1e1cf41e031f79ad75e8902aba760a82728a", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ece306ad1d53729634a1a4b888cb9920ec52ed543dab69591bcd0174b01f81c6", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "a6233c791ed6abb4bf7b5775266f644931cdaac9034329fc18cc28414ecfd713", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "4a6cc84468c1520458080a2d64509964b71128bae10761134387a67a7c23bb5e", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "e56829a8200014832575f4c6d32ed94d6a4a19a6df8ca08793a76ed1dd3604b3", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "d2590a72a3d993bdf7529497d2119135b4e0af133dcf5faefaddb392ec134002", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "9d0a013b3fcf00b6b2ffbc582b73eba31966139075054921e990d450c8dcf7f5", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "bf4fc8d7bac689920ad987b5a30eb6fb3f93a71587054746d76c51bef05cf551", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "64197a766ceaf9ea5533bf0b404f95a6fcbbbb2ee8ff54edd6952b71b985ffab", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "6846067a2d1487e05024742dd58b8e5172f5542272d4ef8f0773936f8eb9307d", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "25aa58677f3e36c1448dd3c1639d5b368df0615a7bafbe73e1dfd8b99049bb57", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "0d5feb89f1485ea05baa6d0720b2642ee543a8a21ab1dee696f65c7e70bc000d", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "3bc9ee44bfc2c23cc4f0ec9e03696100b4165951308e77f74731ad10a1bf4eab", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "11d4dd5ee250db4f6ce6742d5259be666abce91121f1715fcbeb3e8ff56e5610", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "4aa46cf5168063b889fcd110b69f6c19f8b82386ad5070a65b6eb4f208e5dbf9", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d6f61a3c62a2bf5ed3ffb8384dafbeb7b109eb791419ee1dca587fcf27f84de7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "7149d07d1d5a99eb1e1b784b34fd0a5a2bc867db692966acf4f7dad77ebd4653", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "ad9e493f0df21609cb1bc3ab5b2aadaaed8bf09bbc7c5659ef70873b3e3a55b4", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "3fc0776b67258bce7908e8b33fd25cb7b68abdca6ff0d49560292fa6b66205da", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "4f72afebeba7b9ac803b07068f7f204f25ad3e477bfb93a1bc7d78e3ee2485be", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "74b64e5d0059acbc8c8331e7f6af94b208f732200972ca917f0c79baaf337048", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f3fd9703455e96abe22998e496c542335ac7b79b54db61149f9fd5cab4b09b75", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "82c1e527ee190a670f53edeea89c221e94142ff5db7df70be8fc20e0097ad57b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "4ae6702fcc7aa67d8c11f4f70aab9880672932cea6f0d3ac1e93547e5683ff74", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "7d17fc69f55ca916f370aafc38b97be6464a62d68ff38e2c437a883179398868", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "f95a7da2cd3476f769bc01ee8759e6ffc3ac18fa91a52e6f03a86fac8305a9fa", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "fd6ba74540aacfa446959c26787870cda1d42e9a161ac5aa72dad751ab0c15db", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a7e7fdb4b8aaeb703c42f2c6dd16aaa714d8def22972c259832519ba07d3ff50", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "2c7651d7209527ae48688c9da50f00d72f8571f6ea5e5a64298ed90808241da5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "51135f2d161971e27496c297d5f1643687198248f030c186aa0a6a9213d150bc", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "cacaaef56a77c9a6008b4be342c38b2ad25187b7f770905454d411ee4e1b6888", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "786acc368172b999997cc79c1b623fe38e1be170950c15ac3073de083aa09c34", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "5aa1ff3a893518de87a57e24400266d6453cc4d1d5e881e7f77bd4a03bd47647", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "009f14c4ae6c7e7ad22e379bc952a46391593e1f99da69c8887dcb31e6df129a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "29accd72f080d250a8041ff2a7743882cc5bf23509ed2cb3c8a035f31cabff43", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d9102badc7bdbf85fad69e6241e1324b274642e6450f6a0f4342db6dbeac4da9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "df5c5bbafc042b6617feb11175e75eadb2b06359f6189e3302365744c00d19a4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "d2bce3c20aaf10c5c0d686bbdc3cd5741d6292ce8c57b527ca45ce448c81a15a", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "6912627a2d7dd4e6db651a4727f03fc3e208b156cbb2eb11f58186aa147efaf7", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "c03f3ac294988a9f314d7847760fa19c967f8ba1695c2571b9cc7ffee1f4690b", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "2a41d71374307c971acf1a552fd98a12015c2ecbfe068f22b24d2a2290d3dd8c", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d002a38d9554f60953f27df772af37c084a69ceef6fd506f1bbcb0666508cf4e", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "bcd62145d41dacdd4720e34a44f9e4ac9c5a60d06e7a1ec2ee38109c5b4a448e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "36f6a03292bd94a30581c877acefc9846fbf84a8aaddde0fb72b2157df2bcc76", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "125127114441040476a43c3a096fea06b7168818fc1d3c524b86615e19acdfb3", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "0acc608250dd7ad95fa6010c7a1960720aa56f7dbe4c344a035a78afccc006db", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "a4a3f0896f0bbb5c7768e29e64a84942c6fbde63365b6331ecf1b5e2f5d2d37b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "37e0191279b1842eb877b04a4380b79a6b43f13b5a33f259df49138b00af9c29", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "15a0f6401c85f6705f0332254e0dbdff5aebd2566ab8b820a6898f81f37e87b1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "86fd72ad4c95959404d00be026068d8a63b4039106c48b4b93064168e0967325", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "5e18b08fcdc42fc9c7a0a95b9c00fa0f678d8dc9a0e32d2f2c957d7438998a27", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "964684b44df729f5a9e09ce33196715dbeaf041eea4a2464411f661e3bac20b7", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "8f641c52bbf73f22aa4414b341a5a231d63124668ccf979f9894aaef6eb8efd9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "11aac79c6f06412da81877d49195b0c719dc3339f6fdac37e6c6aeda4d28aa50", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "cac3d302bfe19b07e18a956313504ae87392780707af7a583384acf3600ecc87", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "73051ac93a478fe8092c62081cd42a324b9ebd992e7d2f712109e978515110c8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "95d5ec552e475ac2938f401a5df55d98324d2b68a37de4dd916ba7721027a2f0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "70e00024e8169896a5bfcdd0beae32e9f3e5f362e367660504c32ccf4f882878", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "e41e1f3c16b02911bbd3e640cdd6ac8769d0e99fb042cb14ee81e282e8b986a2", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "7f3dc74599db844fef4f0694f842a79403705852861d0396e8d06c1f979ae549", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "8320f7327aa2f872978f779a3c00bd91a708c87691a414654591aac50642a60b", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "7c6a3598b2166638cab272a26b8d95e7f80bf60782e980172d106c339a207ca2", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "23dfc0ade262e9b3b789dec724b7195b6d0c02cb43e93217460d260c555cb0b5", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "d0b1cf58f5fd6183e01d63c0592265e330e0c428a8b74ef4007b186d65e37754", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "5603ead25c168a02d62c693447536f54b29fb23eb01e3b2adcaca0e90228f8e3", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "8de62074fac517c7ba8849df6f98d5479cdbcd1e0fa268d32683e62e7fca1fc9", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "00db2a4b22bbbb53b8f6600387aad88249581c9cac8be93f13ef5c20a43f99cd", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "eff50be490fe63153da534638a16d83069273953da94cc491b2acd8b3ad926e0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "275a3ca9371dd66266b7fd0fb0ef885ae5e10a53f1fecfb8d9a3c81ff468e7be", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "379c92d567a430d6220832f4ef294691b04caa393500881b818373ad2069a54a", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b970a38eac02e160ad285ac2ceeaec311fe61ff9800e929d95643bbed305ae27", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "c49f26765913e9812606fe841336a4d3f85ef3091adf717ba921b8d07b6f53ea", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "26f4a2b823d1a4929c38ef48632d642721a55b5bb7989f15bcb600c9fa886ece", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "6adcb808f9c29ebc3c3726b7d25758fc1fe90243ce9604d4075e32b68a5c6827", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "fd6469df6300e128af6f9d19909b7eefacf6a8c7d3558cafd54773ec72859f4e", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "34792005a5a2c9808ed10f3fe5a492cafb769abcc70a27f87f17f2c12e585add", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "38c51f9fa23196d9d534dcaf708a63373505cfb6b1ae7c605e3786032ca85ac5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "9812f91600d257d060ecb4e1e74346132907249e78696df505978aa59af94b7b", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "448d7dbfa11497e4780d6da6bcbff827ea2cd49732521e25aa8d22a6ebe840f1", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "32ef9944f5ff2aba45dd45880226fd4541c335678e7dd9affff3ed141216358f", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "2a4102b4bd997e47fab3b8fbb8b4a418983d31ceb8697b1cdc13b9cae60889e2", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "24c73f38cf89c0617eaaeab6ea35d0f5a0b3d8bd25d50fa5b93c2d9810c3a4e0", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "5fecc11a29b8064ec7af8e8c9b02d6e9459c4998c39078ea2609d7c8a93cf2c4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "5c7a069b2c36cdd7780254ab544f548fe8d12596668a99c4dc44278e28ac178d", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "8943f9b6de7563c5ca4bb6436d329e2c4baada9176961723b9753bd3c2fe0737", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "cc237158e550cde3a0b43935aed4947cedf408545b992dfb5c79cc32d5b99287", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e42c0c629ce4f69c9286c573596549a6cd5bcea0147390ed4963020f3c641bca", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "d9bcda79a9d1e3d5452d6d65f11cada19a5590f80d25bd082255083921faf4cc", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "5e42384b89211ebc0b523c2a2fe5e3acaf865267280ff26d32d836fde59fd0d8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "4a135adec2db5822eee64e84c0353f2ae9af7dc388e2846de41b61feb639542a", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "b88543f573f68eaf4fb245644c440f09ec1800cbf8a01f629f72193fb69ba81f", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "64b652f1855d2fb85f6343924e677a761ee39867f34b3fc7c2689a69b263ca4a", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "cdd37b4cf5e3bcd14af65597907799589781e264e5f4c38ae589716bca361760", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e39a200e738580456076915df7a45f87a8568a31c286c2580fec83bef8dea08d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "78edbf9f7b7eb349cc105d8993478a45580228a6b6fc3aad0250954cd712627d", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "e01b7598bba97a3f14416b4961404e79f51421383300b2fe0d0a8b135407a517", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "a7af37455033cb5b1206e548d457c69b2c1151a6c06127e35873d8ce36d1c05c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "0dd94c26169aa0601310e2155b754d7c44d413702c75c28b5025935adad3a32a", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "28400c32a167f0970ac1067d4651684f2cff7f68a254c81b5d2db8b4840072b3", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "87d0db0abdf8dbb71a165b58372427f1a86b96bfe7a823346cc582c394210152", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "8e576cdb3ea5120848f7ed5a0adcb5e769b5fb6e74684bf67f6d88374886d296", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "0505bdf3e11a28bb69793942e2ec7fc20fbc0b69458f8d2f7edf20428bc366a2", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "208bf343478218beb4ed8a8374d1f816ed710788a30322b0d820d387ed65e850", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "b68da3e473ec82fcf2d6ab0b63c5787e56a610348a40104a37d5e1a532d83f7a", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "db4b2fb97ad95fdc8d0f930ba98d0b8a1a9a8573a9919fffa7ec17e8b95a296e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "8e94bbf67075782946a29296eadb7a08b397ce8ceae0fa23aa71746088327915", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "268e2804403dbe88522d7e01bf162dfba03ef6b8aac75ddea2476204468e7e49", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a67490fac8b38961812298ea73580d8ec952bfd8fd5f228cdf7ec00f764fb8c3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "8b498dea3d8e3dc381704cbdcb6a69f57418cf61fad4e27968efdee88d107bc5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "c91f34830b14835caa8aa81d22eba4fa9d296deecedc8313a11db52303e48ee3", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "7f9f757619e25e3f33a74222088ac89c81dd0d04fcc66206101a14b618ee2bcc", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "2c53e1e88a4e7c1c704843301f32100704ee96a6e7cb695c572948c89a8df890", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "287e0a07a3801e6604773303978eb7f1347244e4c3953e40f84d60f0a5a29354", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "2cb6fbbd08b10ad9837f39b1bc7863a10e9f7574a57029f86499a018dd784dce", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "cc395edd3c3fef064be381233b9f058857ac3abdade4bad5c54da00f0075b764", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "e5cf985b02c01448fa11bac3a2618729ea59503692389b841fd400bf2e0b0155", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "814dd9da78e8fa50a213ebe6cdd37f3c3dc56234ebd1f9c9004d9c4d504e7f40", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "56150f1455f7aec4082ec7938f2f8161a3aa07320b1285b273113871face7baf", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "eab772f2e92d4934f774f6fe36e5fd79532e9c495c88f1bbafb730c998aedb49" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "5975e49b39122fec9983acb172dda693a902ddf574750920b74e56ef8e67f4dc", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "102cd770d0e9bfa6b1fb24813f6623b84e8f38ca", + "source_training_step": 42000, + "training_contract_sha256": "5cf2baf317839f07eefd2ed65b785e5dabf16971c1df797972e4a2ab07558168" + } +} diff --git a/models/recovery1/0/_CHECKPOINT_METADATA b/models/recovery1/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..8dad30dd7ec113d8f11e36f0ab7d37d41f91a657 --- /dev/null +++ b/models/recovery1/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786162957331139850, "commit_timestamp_nsecs": 1786162957766429535, "custom_metadata": {}} \ No newline at end of file diff --git a/models/recovery1/0/default/_METADATA b/models/recovery1/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/recovery1/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/recovery1/0/default/_sharding b/models/recovery1/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/recovery1/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/recovery1/0/default/array_metadatas/process_0 b/models/recovery1/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/recovery1/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/recovery1/0/default/d/d892c60f4654593696558e2196fb3f0f b/models/recovery1/0/default/d/d892c60f4654593696558e2196fb3f0f new file mode 100644 index 0000000000000000000000000000000000000000..4292932d3cdd609974b9282f9657f0131a336d9b --- /dev/null +++ b/models/recovery1/0/default/d/d892c60f4654593696558e2196fb3f0f @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:79dedd84deeda7d7077e1779f02edcb33a0b7d512257f23ad4f363dc6e0be8cb +size 112929 diff --git a/models/recovery1/0/default/manifest.ocdbt b/models/recovery1/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..37f747edf66a8a37a903b7867fd226c332e109fa Binary files /dev/null and b/models/recovery1/0/default/manifest.ocdbt differ diff --git a/models/recovery1/0/default/ocdbt.process_0/d/30c5f0e29e4f4e5045850e4ae8ac8c14 b/models/recovery1/0/default/ocdbt.process_0/d/30c5f0e29e4f4e5045850e4ae8ac8c14 new file mode 100644 index 0000000000000000000000000000000000000000..09e5c091ceda0040d965d1328f2821f2a5ae4a43 --- /dev/null +++ b/models/recovery1/0/default/ocdbt.process_0/d/30c5f0e29e4f4e5045850e4ae8ac8c14 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:014e464927e5fa53f01584e5b7d59c0390f00060e942334aa49c873e4b0bfe88 +size 447942656 diff --git a/models/recovery1/0/default/ocdbt.process_0/d/5918aa7c6bcacef5d5a783faa7189dec b/models/recovery1/0/default/ocdbt.process_0/d/5918aa7c6bcacef5d5a783faa7189dec new file mode 100644 index 0000000000000000000000000000000000000000..e0050acd3d8dad9704305a66a6bcab190639b2f8 Binary files /dev/null and b/models/recovery1/0/default/ocdbt.process_0/d/5918aa7c6bcacef5d5a783faa7189dec differ diff --git a/models/recovery1/0/default/ocdbt.process_0/d/7c0764d7c8f8db23a57c6d278af3966f b/models/recovery1/0/default/ocdbt.process_0/d/7c0764d7c8f8db23a57c6d278af3966f new file mode 100644 index 0000000000000000000000000000000000000000..2663ee0c5894f1dd53a9945b59bfa369350cec5b --- /dev/null +++ b/models/recovery1/0/default/ocdbt.process_0/d/7c0764d7c8f8db23a57c6d278af3966f @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8d3d183c041c89ec2b1098ea2a689ab3ba56547c1c57eba9bcaf11f1d51db15f +size 10940416 diff --git a/models/recovery1/0/default/ocdbt.process_0/d/ae2fbcfef5c278ec27aa5ca43aa085ad b/models/recovery1/0/default/ocdbt.process_0/d/ae2fbcfef5c278ec27aa5ca43aa085ad new file mode 100644 index 0000000000000000000000000000000000000000..c6a83495705018b488814a0f58705f4373e3d3be --- /dev/null +++ b/models/recovery1/0/default/ocdbt.process_0/d/ae2fbcfef5c278ec27aa5ca43aa085ad @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:266f1e1e48be0aba5fd10fc7496a8dd9f9c101a6d33fa33cf066ea14555b45d8 +size 120410112 diff --git a/models/recovery1/0/default/ocdbt.process_0/d/c54c81baeec72d47ab8c23e6dfcd33d6 b/models/recovery1/0/default/ocdbt.process_0/d/c54c81baeec72d47ab8c23e6dfcd33d6 new file mode 100644 index 0000000000000000000000000000000000000000..e1c7eafc49ff8feb32d13d911fbda5189f3eea7f Binary files /dev/null and b/models/recovery1/0/default/ocdbt.process_0/d/c54c81baeec72d47ab8c23e6dfcd33d6 differ diff --git a/models/recovery1/0/default/ocdbt.process_0/manifest.ocdbt b/models/recovery1/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..22959a845354b858517b4390782545dff55d9876 Binary files /dev/null and b/models/recovery1/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/recovery1/checkpoint_manifest.json b/models/recovery1/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..7f9249cc953c7b1d77b82852809bd671533c33f1 --- /dev/null +++ b/models/recovery1/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 37000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "8dd50b827772d5fe8e3959cfe53916e52f007c3f7e247cb8da9a8ca9c175e680", + "runtime_config_sha256": "8dd50b827772d5fe8e3959cfe53916e52f007c3f7e247cb8da9a8ca9c175e680", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "8dd50b827772d5fe8e3959cfe53916e52f007c3f7e247cb8da9a8ca9c175e680", + "schema_version": 1, + "training_contract_sha256": "2eeb2233a95c5cd06c5bac2cd88c44fdd2cac1789c6968f6477d5409bcd604ed" + }, + "sha256": "00d6b3b4e89f4a7d52c3c5c978d69e1af939c72423489c50ddd7b8b70486c794" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "33e0ec37736c874ab2ad8820cb2743acd25dcdd64dd946789c5742d00e5e188f", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "9bfff10cbe13f0726eef0a997368ceb5a1a8d3e4cf47187d4d963d01a2cdb7d9", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "e9a6f9a8d4c12aaae4bd185226b85065be2effe4bb66773eee6165bf08a6bf4f", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "2887a88d9b0577f772ff82add62135a682634e6e0c6127dd1913a275c01a3e42", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "989f58ffde9846bbd9768ed3dae0b754b3ce7274686fe767dd3e931b61c16216", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "ec93373c6b7f37ac8379e2dca3537b3f8745422d5ea643ccfc6f34afa31892d5", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "728c260ffbca924361ef2ec4a956086673a353062dcf0072fbeae72c43f6554f", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "e09b1ce90dd2b6cdb3f61bd8cb4afa55b66f97adb3bcdd041059a57c0d6aba8d", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "fd35c3be342b38e76d98b7b920684b25ca6ee927fea6365469d3b1f3a3488b2d", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "9de50601ea203f91b81bfd14b276c7e0825fff4874739990082a43ce813b21e4", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "f6c4791d62febd38dbee8c19dc40ba7dee634bf7f387984a34f92f79fc465694", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "73521504a9ff89db37e1a2fbe33d654874f18952e90884d4d5b73d0794fedeca", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "99c4f1244b89662e05cc57dd337d1806df81ff3551899645de7894981ac8fbae", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "1e2ad3b041f8411321161d2ba5c7aacf21ce32372e040c649217d979ffc5a439", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "cb1fdbb257995bd055440010469cbfd9f339bf4f750b368d17578331c20ace30", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "23b68a134152bf4f95ac2f52aa2a8b6139998283b62de82de088128367bb201a", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "7f556ebbb5a77bf0bc445c9c95511ce1203d2f0693943312669be34089ea6e42", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "ade9d10049c08ea7457438103f7a5e3682c475fff218a3187dafe6e90ba094b6", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "492cf7c0cce71d02e139ccb5282f782f94a93e69308e5cee793cb38b2d7662f0", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "730a01877d185ad2911b9825e5bbf7f3cd39a27f0591712e4d3490ac122cb836", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "b9d23d4cc7c1c86c2c5c0aa629abfb3c65349404961f645733093d145bc6c67a", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f23e63dad1f726e984fc52e93a8732275bccd33455c5a924c2cd5f03ed757e8c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "dbd94ee7b05884c4316d06541c703aafbd104bd86ff4e77cb357c7140a6ae93b", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "e67f53e3ccf7da22d0a415f053caa5f233bc2d4d3f4a3d43049e8493a2cb9b6c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "f6cbc46f1ad11856cf61acc92c207fe31326ad29d388fdf63a7ae698e6283614", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "d2d06eb2bf79c087354e970750d590d559ce9e84f1449370a111b25b7e241ca0", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "a252d92ed45bc734a8893f6eccd0c2c3a357a453f30ed06697d6bf7673a71839", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1626ddfb57d3500a3f361d011f6041a114bf6d2fd8ceb2fc6dca6bd5de890acf", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "250d8945bafe960e8fe4e20999a90b862f93b8541da42baa0b4a1ce72b8e0b73", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "30d4e672c9356a838ba54938f66b2e6a79f3bb922ffaf933aa5fbc3479dc82f9", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "34d4e5e8e8617d9e6b9a1070ac0111b4d7c04e2f56a0a139ae754e905b9e303a", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "eda9bb1c8ddf5a30da65237e427f4d550055ebf5166cbc8ac7c3f0c82bfcc47f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "b1bd79ef42c5629602b443db858d1966701be3b91c989ff1bde7adc97cd34738", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "16b32725374062a03a96d01e8bb327320da1aa76317703bb6526c4c13183bf6f", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "160b61a4a131db91df1224d3590f4ac0d57d461f32cbff12520ac22fa4c68926", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "1e569da34ad0082489069e6c2a647ae189f06cb5e27c175cbcfa02e98015fc54", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "d6a3e703a7fb5250397a6364334b1ebe0f9489ba6bff2501957e3cced87038fa", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "02f7a63ef668c92d005b02a948e3d1170bf261b971e398516331405afcaece7b", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "254268672242382afc63df185d0ef325a95317d90b92b598e210e0b6ceed79ad", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "c3633449206e12f95ded6244be88b63e4f65a1115381ba12f671e2176362f404", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "3fabf3f9b0dd93cd2c300bdcb43d3688bf9896cae08981d5d6a139cda72ee8fe", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "746e1d24460b85c6e039ec55f2443eca767c2908aeff46e6b42308c9fc02c24d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "eae72914285e3fa4e0c86a31144c1580a04c158add1906478421ed01ef72627b", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "00275c0fcaaca4cf9c3970da14c3c1a169dd8f912c5e8e97ff1efe2b37a48bab", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "4e91d77446649890efc39ee1464d9d97dfebe51cde1d90dd72b9a6c3f6ff7d60", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "5dca3bff2af41fe2feb90541f08fc1e55fc81e5107d3e7cac7970dbb4f69f5c2", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "be37ea1154070c9cecdf693d222c4a5a0b0487cc28dea2ef779e7d8b3c26d74f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "190fcf7681f4646e1739cef574662b073168c0486adbd1f333e72db98c9bae4d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "5faf9a8d8c5685f99b3c2b2b27ce271a2ea66aa317fc78c2822fa4d8dbe24eae", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "7156e92287f27efe0ec048d4efaa481c22f5e83bcb217e8cc6d465f9caf92b2f", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "cdd55403369072f2328c8b42db5acd644bbbc2dcf8ef07533809e4c8a7e7d326", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "4424d00a4ef0b9363772af9daec4832021b790c2d0540957ba2f698a325b571b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "35c9ae9c77f6dc8ec42b6a018aa5f4eab77b9e25f5f989eb510891e142dbc006", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "268064b418e7f63fa012c672e698c92d9d06d77408e18ff159b5c6282191cefd", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "af47a7ae59fe3a8c20931aaed462c6dfff61a9cc1f9547760cdae9fbdd2cf2eb", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "aeeaf730fb879eb83a95a78bedd5cafa60e21f78b5e6a32f1d25a51638cde8c8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "623f23b96b0f28e76b534afecdd9016a559489caa7931cac6e5d4062835f64da", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "5f85cd0c88fe1a2467fdd28e7cc20f26b3fc970ffc13cf0146e032f35d2cda1e", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "cb5805524802078fe97a035f6d81253cbc749c608ffb9798f8622f0eaec1b82a", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "9c9195709fabafef277a9d769d2b264e5f1896beed0feaaa8c35d582e5c3b03f", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d2c744082e4908bba09de83d3216401ecb974bf6ce6e40d9e37d53329433f508", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "1707ec8a71be96b7c86fdf73ccb8d55588a1e45cc8c97f00eabed1205cf8badc", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "55f3eed66ad5593ef53529471b614cccb3c200adef87dbbee3b4c94274fdee6a", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "f391bf81c99a57b040f69dbe4df080c1954ed83a4cb25d0eac5f4d9de202d900", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "1248a1fcc583fd37e7501cf88f93fa5daffba645eda42454eea23ab746a611ee", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "26a3e0024bd653cb11b81e9e2de5c8242f75fe52a0af92d89148f659a406d543", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "51ccc18200f6ea75c902054f7cdbd1bbd042c6f64bdd2a5b505d53ed519306dd", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "ef5a249c4ca121945d2ab231dcb50514a438e2d67967ae65cdbe24ef51c6101b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "74cc1786c451f92562cb06eb9e8da89830db9834ba274ada5e3c9db5f03e7b74", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "567b03e79684431f870efbc4dcee22715240817c79846e1ed30a7667892c79e0", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "3d0abfd602a288eafae393edfe471e3dcc221c8d3d45c640df32157817c3cb04", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b442a5da0a2e8a7566c7081fedb545d0355d5db81386cbce085f61e12cd96bc5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "9f2e9f483b7758442daa4e3f35d1e070410fc78adf6eca6f2d5b4654eb08bc31", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "5150bd46d4360f72c8e9828a244a71c3f38f2f55ed1c1aeef2ff279204177560", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "c3fd9c379df0358e07b43261bd58c4426e7e37a65bd83d644fc2c4cf2b1d439b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "699afd1471d3e29c99c6fa38f377b9d9b81ad2ecdc20850086acbe19227c8e4f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "b4f48d44a773add32dd8b37dcf3005221197c5712ab1d7a2cd71bcac0a753a2e", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "4c038a6398cb6e7e1c77f082d567fb0081394a9e87188e4102958fce086ffceb", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "de3c8e000b679b11653a102f1922c5711b950d44dfdd15854f6f12b0265ce9c4", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "124710805ad15c69309441e308685a712adbfb4d98f29a2637de35328d7f211f", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "e54ab066a69735e82b3fb3040b2b8c560ee638a324703b8cc466d1981206f171", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "91f3670b2e31fbd8a3f3098e92539413a92ca160023a349d426af5fb7bd26227", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "63d37a7884f0487ae45cac68331603c814cd2f9e35ddbf939e22695b7e4177c6", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "cfaece8e139b6b02b2d5b7ac3fbdc7ed956d2ecdb4cad1b5d5e30e06bc66e8d0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "275dc0740fb1541334e90b5e5b4e1111fd47071e72067dd04e16f4df669d1f7d", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "44496fc25b085ea7d2ad8776f9ca95567d05e7dfeac59209b61043737bbbe073", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b598d00e2ba92863827235f410ff8b7144b063ad7a7e2e9eb2b4d007eb829b9b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "89b9be1a579b222a79643f6cd292e64969e16cd7d4a33168e23bf14e2f8f33ec", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "e25b9061d57cdc9864f8b6fbfe11e7fdebf05ffdff52c30a2e615feae25294c4", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "32cd37bc1228aba6e9a1ceee681fbf114a397458c4f3f5b1cb63cd0df3c13667", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "aca2b446654cbcd18e9117f7f11f8085c7b19054cf1b441da869acf0b80ca98e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0fcd5811ac861828b93cfa08088b249c42b771018b26e5d75df11d747df8ac2d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "b4e99e8e9e707c1077ade36854cbbd6091c5da1a1174af46886675fbe143fe77", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "a8a66c4fd25e2f4138409ae454cb6fefce6686ab9516492e626cc8f8ef8f5faf", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "45085ef7f90043d5322bb5ead2a65b4803c12737c0d767b1c5bad93a618aff23", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "171e77ae956f5521787ee8abcb30dfe87ac80f55f9e9bd2934182e45b24b3904", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "0d453b7ddcc277645c13a1669722eb7c19bde076d37f3a64a5cece2e7348d5f1", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "570926a30ca02c3a22fc7ee1bb77f0264a8224cd3c31647ed17a8486de932705", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "c3b6b0833b77254f72313289518503862bb3e7a82b0ed1f199e26a2ae9963d67", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "eb2b9a54430701dec2b4e12a0691e249e12ee53cf890ae5b918486d99efc7367", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "594c81e13a33d38c510ece343ad407732d70c3268392b4d86b4bfa7e5aaad5c0", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "725bf891d5fb7ab9a028c486d0a9aeab2afd72b7035b882ce0eb7a2d7069473c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "2fbd82b64c152f583eb9ec4880b1cea46e4d134d255ab6f697abaf53cfe7a0cb", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d9e511dc56512beac5fab2499e5ce55cdafb6312258813ab7e6fb85b25eb460b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "d0ea9a96a3a83329d37424f4fdaee3f9ca5eea7c76fa0e0ba34a1f3161b3a741", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "10689ff89328578b3d73deb0b224941fe84ce93fe483d40782e88f9fc82daf2b", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "19de3b65e6ed52a41b2a10a8712cc5c7fb7f6eaa4275858329f3b5b7024a4984", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "6964de66ebe4ed0ac0b3c73d75b9d31aa38077d415be20de6a468f8b187305a0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "94189ae30d6a71f24027745a093b4d9ed4919c231a2c9bc47d2e12f7811fe744", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "73cb1b73d709532873d35b3a35cf3f268c3bc6856e38642b8e43af66280940eb", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "24a6944f1671efb11db16885ea6926c01ef43d986c022880250f804c08992794", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0a0910fbef065d13db1a4e5758eb2235aeccc3578f44846a0b13e18b940291fa", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "1143f222e6daee488db77aba7e176e1ac2608999099c7b38fe98fd2f8a54a94b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "05fbd90723af16c80579df6c3947629644ea5c31d874ee8e105f5fac0391147e", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "35e1378e0a03f882e533652d66539672ac96714bd352b6b67b0484b16d933393", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "3e736ef05d3893e03e0b0ee1958f7d98ed6919ff89cdc468a2ae6192a8baa14d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "1266c094a1303643f990020d9130b8009b86c697eff16358cc65a4228c5134da", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "3f59c692a12ba4a040358f8fb57fa2002aee19970296da41941157a646fce97e", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "bce94b88a73016fd3dd2568ad2a9d46923d55f7f4b96f54220124e4f910f5a35", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "4c4cc06dc3452b0c435e62b6b259edb1d11c739f9ca991ac81134c34a7beba15", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "bd558a0722bef8c88924b188d89e0bd335e518ec395728ea17acdfab395e8809", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "42e855ba3e9d45abd55fc60a46e444755654f7ab6873330d6f5bd0869c2ce9b9", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "a48703763c3ce261a27422775876cb44eac39cf978c361ca100458530cdc91af", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "d7e6c4cd255d4e0778d1fd0d6d9b0f1fb72405f04a5ca7ecae9065e856c64b41", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "3e4eb4d1eec255a77214bd3a2043a5834b635647083252444441b4474d376b77", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "fdcc63395c0426781e7e7fe1fb6c71b695c9ecba18f8beea3fa4cc54a21bb244", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "57ce3aa3f6b63e128c4d3b826886c652703c4e4f250f4469840a0fa14973670c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "3bb99274e2ddb6a6554ad30c6f833709672ba80322f8abf4461ab79e4025d4da", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "d109d65894c28b9c8cd40ba2718e88fdc90d2ac07dd73bad26bae5af15f6bb07", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "611673eaa2188cfd8f486e4f0a83dcf8f10d4d0ae2a2b88b30446c62957be188", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "4eb6b5e283e0c536c80b6845b067f95b68523289ab4c9b99883e2900056da794", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0265a7858e753f9aa1df41f97b0ab4becaf1c2640ce4aa2fa3389cceaa26180a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "f69822fe6a4dbcf3b6acf048f90ae2374884f9fa5aef8556c16f97ce27c8e309", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "f66ab5b087ec887f496e8879806ea9f9e72a5f18b7d62e87bc1845928dd19223", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "90a197b177746c723358f55f6d08dc69a14e09bdf32faee005903f10d45b6ad8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "20b641e5fab629a866e6729410db02bf64d25c93efa0e4f83badc245227b30d3", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "bbafb37a9793dd27f731ad81da8f1e68c3801886949e82b2f4ca2d6feef69975", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "ceef91f2ae838e54c8ea5751df9ba1bf046ae3cda088b630dbb780f0d97ac49b", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "cba2cf2c674a51ea0426643ec570c7a15bec86e9e0f1da255fafc84fd5d3239a", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "3fc7db63c91c182a4588657a11878d7eaff558b08e3e4b659ae58b8a05c83929", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "c7157e92a5caea03a8353cbcb760bb4f0a4edfa8bb301f752198ab1a168cc1df", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "694e3864ffccc41c72ad2b69f5e7257f33078e7bb2db5a9a3c4ce6d245b15f61", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "574c905fd9e7785993b2fb96e628078946eac45ffbbddc9d6fcb23f4e23b627e", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "a169af1becc885bbcc6034a2ff9d3544b3cc58fbb25e20f7ecd98dea167a5f69", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "a89fec864836918564e581b3efc2b4c964575b450b6561e0487b15fb6e4e3990", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "9e9cfd339f145bb91b8a1cfa698d1b3397c0c28eadde9d22303a479db41e307f", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "b96c9724404bafdc2970f2f83afa7a81a61058a930f39348545947c4cedcb702", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "194656b348076f439bdd6b9b35ca9eddc4579c3675afc9cdab306ef0fdebe5f1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "bc34625204302da97c76e8a6b51241a47bf7ee0527eeea16f63868ff837dec47", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "c18cd0d397f416fdb59afd07a404bf1c47d8adec2b21f6b8feab1d642659f8eb", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "4ebd90ca852884abf69e00f05af55c66b2a1492779222e3c6bd5fc2984922bfa", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "477f17071baae7c33afe4fa24a5a5878b86c12ad014e4ed354193755bf0cdeb0", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "7eaabb3d613422e2e5317d3ab7397e38e368f004608e17bba3ad751e29bdc7b7", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "1d5263cb49769197cc9814b43da75b7258fbcf22f3e0ddf848e882b365559d8f", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "64316f063d5140ee873bebb07e20a1520b3f54cf5f186fcf4b4e5859c5e1c844", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "55472ccedf8d28ee4b14a0f7b9996bba52f238f4f9f56bd387c61621092cbc3c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "a77fdd6c37ea2e8f75b1a48458020082650e52264a2283cc76f8bede4ccbac60", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "67c92bfa5a5600d80e029a180d932c14665f4c58392a82fa409f248b2453d9db", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "fde0a0eed064d62de0ef99889e457dd5b2552469a2b28dad9298d4b35355ef4b", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "80025294202631c6e75d0a2e95ddf0ab6a8857a0975b5d1d29b19ee67f894344", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "021c089a461d19d98a46ea8c853f362e1e2ba71ae5b62170bfd44973067f6e80", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "5bead8c099c74807df921585185b749e206ce38dd7d848802768079b321c79f4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "bb272de152449530182c44db5ab2f2c2ce579c6f2371d3a3a96a69b9efe78dc5", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "816fa9d911c8aad44b28e2f72e2eaff9c2648805f414b6f427d02eba2707b446", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "dc291107e410a254b714378ca8a631e1e81a079def8ab9559f4e1c8faa02dda9", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "3ad1987d468918610a765fcf0241730bcd21836f3bd8f375c0d7a388726fc02a", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "48d9af338001e3259a9447e07056d42f86501a3ba9beedec1fbf0b27ec96f4d3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "1f81cfb9d430ba4fc7dd91841450b7fa71f55b051e3861e1f532398920dca1df", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "0b9f19a6dc5d5649f95c69f289652c2601f0cfb5ea67de5589d2d7f97f7de618", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "45fa0365db437a878999cbe288f561ed7692012c7da15594d60b0f6b1237ed22", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "0791a259f711e1b3142a399b5e1c2a653a2a63c5de93e0f0cc9c4e129aee3186", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b679791141989259b16839dd06b7ee80cdc3d129e4b73d8b6cc11c749f13170c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "1379496efcf4213bec8fa81ece713c40f0a0af6cae3cac771f4f600749204b4b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "0da7b93bbc07fe0bd128c161ff2e750da71a228fe9134c39a80d614c80bc3c02", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "6d632f051c95db89b076c828b33a4061a6d1fafadd1404cad9e85d14eddb9991", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "5340dfe0c96cc98ebd7cc4e4772ab38915cd5f5b9970938e0b58c35822f27e3b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "c61534b04daee2861a24a6c6ab51c125e618962a0573cbdd1016201b8fd1c66b", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "2aeb2df45827af170ac1a64e17944a7615ca59f53cad9276839cd03bbfcfff8f", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "f8fe9ab0b82f998d51be74b794e74256181c820331cbc7c725004db06d443deb", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "60029b70e4a1798182ebe0861d38d61cb22cd9b8eb8b84f49a30d6230812d9fb", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "8b2acf703756ef6f8226deb0731eb7be40334bd1fe042bcaf53cd544ad97c909", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "1e219c520a910d51b39a71441920d899f05c88438da226a05c29a7c709707510", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "4254d110d1ad184d1253b004746a2478e04430ef86aebae1724b0cb8b73e337c", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "005b5c8406633d7da7069b301c7ffb7eaba874961fb610c4b0a16899c4bf06c4", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "3c4c74065bc58bd24e6051bdac83f4a677a5f010ea71ccbc2f1f17d482215d16", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "02b3f7d12951dc68bf4a42598bba893edf7b04301a7913a197af13d6c4f05369", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "992e791cc0e2312bcf65f02f8128b767e241369b8fcb1fe19a5fc57152147fd3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "a7f2496503d7ab7b2a37a37fcc675e987220cba603f9adc2af0074a77c1b6981", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "dfa83bd43d90d399eaac6a564bd269e0b5fbf7c52a28ead84840189bd2a57839", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "407279c0eeb0f19dd02cb90c3e28aa8b83f5d360ec5882dceb4ac836640a186b", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "d6f57d77c0c5d5d8befe8ea7c87369bba0ab9aa82dc2e87cef0983b0a63966c5", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "4f477286ef92522b71f9daee6f3517d03c27a1fccf46e919ca52876ee4f9fc80", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "06b205d905454b028f786102804a1717721b97e1a9e907597fa8f3c7c49f42e4", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "b7d2619634185a6b3cb9b46ebeab9324a9d9055e3a97f3349a2811ddbf6bb0c3", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "eb7b06d2f1f34f86d855597da49fc0ac403493ae96e623d129607a6005c6a2d2", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "50ae05524fd8deea912a34d6d887e40601b46ed125f02222b097ee004094b9db", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "2452a9a96552ba68110940971e0ab811fdb216f6b5f79faec14a0e43ca4cb800", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "649b0468070d94482f068f7d6d8fc8743f1b74fa5c442feb9ff6c13f6f75a25e", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "e88df9550544926a31cfe828323e26c43b15dbc1c37ac0ac8391b52b7afcc0ef", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "1cfa839b1980ca7d293ded5678488018cfde40db76da75460dbbfa5b71f2e370", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "4c8622f8514b4bf8021568c1da48b862842424b9d05a5b51e225587521b6da75", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "c758997f27908cc8e72325c1fe30b143b9d247d4e5505c880e277ad161643d87", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "2c16a73038b9d06bcfd386a4bae8470378732c2b2b85d9503cc309abc894f4bf", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "024c2836ec8d03159c96d072deb3c91f672907acf62558e9a23c1dd83ad577e7", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "7e1f7d04ce454dd1675827bbcc41636669d3f4adc896a45a9fc30c06ddae2549", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "da0658fd83851cab24abdc578a57e5cfdb178994e0e455c12ffb7e258205bfb7", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "e7cfe3d8f06f15039221ec76dcaffa425b9a9b34f00c1b0eeb408d875866ede1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "20c571b23f77e80beda59614c68175d2fe19deda6c3d5dd3897c1002ecb1a951", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "e88b3ece725f495acc021598f1bb05e0f9b136b6850d0fc9eed212251654f013", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "a97a18ba173db7f8e62ce0e8a1bdff544475d83a9d115b87c4b26d5be149d00d", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "fea0a4dc9095b9a0617d4507efc05db0259ee30ac1053cb27040a72295c00fec", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "ac414e31daa1323286a27d3ea09d077590e4e10b1672c0010235c4681f15cdd7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "a7ed4d773ef9fd433ac34fda3d6d285a4acc1c6dc09e5a1043fb8cc747f4f48a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "01b4c8160f178432bd67a23a8231335e60f8ac644e8f8486bdd68ceb2ca9acbf", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "21155d369f8557b8bcf9f4a1aca6d823c14918388d340539633ed519fee49597", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "65bbce12c7baad685b04d37dcbcfc954958e2316d340544c04bb0aca07576994", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "a3a4ad11a43e1991232efb0bf574d2725f38aebbdd58548ef4bdedbdab5a93d6", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "51ac3c5b95db41f2b71a3232db00841aa9c152b706fd415842c9a1bbc7b4b15e", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "ed42110f3b3e40cab2c6379f811033e5b4cd211bce13b49482d29bbd87d641ec", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "fbb92747584bfb8d3093d82b7ed0343047e8fb3778b05ca0d05eeed6cfa5f9df", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d921babcdf484598aa5a99b73fa67d7dcdb9fbd9a988430ab64eba39d07e5645", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "02559c47d9272421d9f009d053d79297cb76663c5e26c6fff286f62a223f5b0a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "6c2ee392a277dfea2ab4416d9fe00b5a1e086711e80fee73f91f49502c7fbbdc", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "5f600a39f471369660d5a5aeb5da55849b0155b1d56c7c61fb6ef8f119fa6e3b", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "6d484dc482929fd7171f716caaab2d09bd5a2db49dbdd4f540338701eee5e045", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "024c8c0a76f79fe6fafe0cac7e9536ef03c2f8cc931614ae58253d4b4636abd8", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "bd3eff8073ac6ec92eed26863f21d38a2114fa79939703a9a9793eb339c13c9c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d3d4443f245df97cd9bc8412ac9d324e1a44dec0846442e3cda6dc7362f01fd2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "6d4f26451bd4e10eb25bcf7fe1ea7c21db71166b15a354a0bda3bcc0e23222cf", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "700b5ff9f0b70d4bcba6539971b20d4b051554542fe5df9d89989c20299befa4", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "4f81b72c456a3a75503e9b61d0fcd435cf61772f329f131311c8b17acc1544c1", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "b4de081eba16e8b0330d18c7ffe94ff04d660974ab4a45ad40e415977c5665ab", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "b5c3b0806e9ed16abf5c85f4240f830b59e81ef330820d33b0509cebdf30465e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "cc8d8aae9517613051826e15f0bb6490a390b0b2afc6702bd84a8c5b019ca060", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "88593db54656837dbe97e98912bcfcc2ec9d61d69f0814f95bdac0d91c9f64b5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "00639640eb1e1e3a3bf05a42645cb71a058b0c5fa3c4448bc8921c94d489cff1", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "932b830c1611c48c64b5932063d256d0c06d044900858b3a65509ec6f26a420d", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "3ede1a6054c34de5d05f2bfdfeab476dd2d03b556c52fd28c163f004a171139e", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "42662ff0d4c142cfa061b489e1ab19e400470cfd11914da1dedcfdf861543c0f", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "bc9ee31f0b161feb342b8916100b37a25b5fa9115abd7bc1e4b7c7856011e7e8", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "5b8126d949d76895bec8bd8ddb16b4b2a3b554f9f29238018fb1d439619ce265", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "afe41e2c942fe19c2635ec0b9c256026b70ee70c456b862d12dd0c559e3838c0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "fc4c1b72d9429b4ab48e5d7525afca7dcc93cea5f04c7ba13ff699620d9690e4", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "4d7789745dd7164f353639431f672cc1f82527b4325c99a0b1e546087b67e8ce", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "440f404e8e2adb73b91a953b8f1ec7743531c4a93f750a3090a3df997e6f7051", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "45556edd31c7d94879d181d0a30af936b3c32e78464099dee8ae465189bf216f", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "562560ca3fc3c761d764804ecdd3b2a3befbd130c5ebe34698d46500be55081c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "78b2a1ff21ca85e6522e113c34b5d130c5661549d39f0cbba1f7bf2cef541f19", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "495bc3baca1440178b4a1e0477b0d2ce1f7194576193c8fde449ac828962f3a1", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "f050f757e83cb94aa3fe81997b8e23f17e8d7f3c89514b8ec0f15006af6d5850", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "6e6378a13aa541823288d61c4b0799e9bb337bba4b4f858489970dfe948d67b9", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "fca112f338ecb578af51b16a8fb015d6a86bedd147499a1228909abe24b172b5", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ada951270659fd6ac33fa78a981934f6a87a47f34d3963fdbd9ba68c4e6e5375", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "fc907cd39aeb9319108ffafa85acf3bce98bd58ef8334c698eaa92f5ba92e2cd", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "fd7ea603923adad50f2acb72d781b2bd154e7b1edc2e2a76b2398ad79ab066aa", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "03331fb04c2a4262b2313549540b1435e09719b68111c3a25a6d3c51d7e46740", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "d3cbe857198394f6192475eac1de60d5499971c1746f1519a6a3db292e17838a", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "61986cc707f8c05d1d116b57bc5d8d544061e2583f59d084a87764b2231931fb" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "ad06bd535cdbe6164ecb7c5838bddb2ab43e49a0ba24201a334f65ee9da98454", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 37000, + "training_contract_sha256": "2eeb2233a95c5cd06c5bac2cd88c44fdd2cac1789c6968f6477d5409bcd604ed" + } +} diff --git a/models/recovery100/0/_CHECKPOINT_METADATA b/models/recovery100/0/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..9e24ab62dc0168382ebacb80efc14ab84b1113a4 --- /dev/null +++ b/models/recovery100/0/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786225532225230916, "commit_timestamp_nsecs": 1786225532888434250, "custom_metadata": {}} \ No newline at end of file diff --git a/models/recovery100/0/default/_METADATA b/models/recovery100/0/default/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ee08954c6382c2f49fca72a7fe221aa4377c4de3 --- /dev/null +++ b/models/recovery100/0/default/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('action_embed', 'action_max')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_max", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'action_min')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "action_min", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('action_embed', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('action_embed', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [126, 3072]}}, "('action_embed', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('action_embed', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('action_embed', 'passive_embed', 'embedding', 'value')": {"key_metadata": [{"key": "action_embed", "key_type": 2}, {"key": "passive_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 768]}}, "('current_latent_patchify', 'patchify_conv', 'bias', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "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]}}, "('current_latent_patchify', 'patchify_conv', 'kernel', 'value')": {"key_metadata": [{"key": "current_latent_patchify", "key_type": 2}, {"key": "patchify_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2, 4, 768]}}, "('loss_scale', 'value')": {"key_metadata": [{"key": "loss_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('next_frame_index', 'value')": {"key_metadata": [{"key": "next_frame_index", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('num_cached_frames', 'value')": {"key_metadata": [{"key": "num_cached_frames", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('output_proj', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('output_proj', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [768, 3072]}}, "('output_proj', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('output_proj', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "output_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 16]}}, "('time_step_mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [3072]}}, "('time_step_mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "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": [42, 3072]}}, "('time_step_mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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]}}, "('time_step_mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "time_step_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '0', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '0', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '0', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '0', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '1', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '1', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '1', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '1', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768]}}, "('transformer', 'layers', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64, 768]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [12, 64]}}, "('transformer', 'layers', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '2', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '2', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '2', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '2', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768]}}, "('transformer', 'layers', '2', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64, 768]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [12, 64]}}, "('transformer', 'layers', '2', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '3', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '3', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '3', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '3', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768]}}, "('transformer', 'layers', '3', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64, 768]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [12, 64]}}, "('transformer', 'layers', '3', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '4', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '4', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '4', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '4', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768]}}, "('transformer', 'layers', '4', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64, 768]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [12, 64]}}, "('transformer', 'layers', '4', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '5', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '5', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '5', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '5', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768]}}, "('transformer', 'layers', '5', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64, 768]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [12, 64]}}, "('transformer', 'layers', '5', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "5", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '6', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '6', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '6', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '6', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768]}}, "('transformer', 'layers', '6', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64, 768]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [12, 64]}}, "('transformer', 'layers', '6', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '7', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '7', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '7', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '7', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768]}}, "('transformer', 'layers', '7', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64, 768]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [12, 64]}}, "('transformer', 'layers', '7', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "7", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '8', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '8', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '8', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '8', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768]}}, "('transformer', 'layers', '8', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64, 768]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [12, 64]}}, "('transformer', 'layers', '8', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "8", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '9', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '9', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '9', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '9', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768]}}, "('transformer', 'layers', '9', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64, 768]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [12, 64]}}, "('transformer', 'layers', '9', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "9", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '10', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '10', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '10', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '10', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768]}}, "('transformer', 'layers', '10', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64, 768]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [12, 64]}}, "('transformer', 'layers', '10', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "10", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'ln1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "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]}}, "('transformer', 'layers', '11', 'ln1', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'ln2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "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]}}, "('transformer', 'layers', '11', 'ln2', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "ln2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('transformer', 'layers', '11', 'mlp', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc1', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc2', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 3072]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "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]}}, "('transformer', 'layers', '11', 'mlp', 'fc3', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc3", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3072, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768]}}, "('transformer', 'layers', '11', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64, 768]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [12, 64]}}, "('transformer', 'layers', '11', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "transformer", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "11", "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": [768, 12, 64]}}, "('vae', 'vae', 'decoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [512]}}, "('vae', 'vae', 'decoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 4, 512]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_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": [128]}}, "('vae', 'vae', 'decoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_norm_out", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3]}}, "('vae', 'vae', 'decoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "conv_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": [3, 3, 128, 3]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_0', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_0", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_1', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_1", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 512, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_2', 'upsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_2", "key_type": 2}, {"key": "upsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'decoder', 'up_blocks_3', 'resnets_2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "up_blocks_3", "key_type": 2}, {"key": "resnets_2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [128]}}, "('vae', 'vae', 'encoder', 'conv_in', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 3, 128]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_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": [512]}}, "('vae', 'vae', 'encoder', 'conv_norm_out', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_norm_out", "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]}}, "('vae', 'vae', 'encoder', 'conv_out', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [8]}}, "('vae', 'vae', 'encoder', 'conv_out', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "conv_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": [3, 3, 512, 8]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_0', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_0", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_1', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_1", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'downsamplers_0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "downsamplers_0", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'conv_shortcut', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv_shortcut", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_2', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_2", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'down_blocks_3', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "down_blocks_3", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'group_norm', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "group_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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'proj_attn', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_0", "key_type": 2}, {"key": "proj_attn", "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, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'attentions_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "attentions_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": [512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_0", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv1', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'conv2', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm1", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'encoder', 'mid_block', 'resnets_1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mid_block", "key_type": 2}, {"key": "resnets_1", "key_type": 2}, {"key": "norm2", "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]}}, "('vae', 'vae', 'post_quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('vae', 'vae', 'post_quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "post_quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 4, 4]}}, "('vae', 'vae', 'quant_conv', 'bias', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "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]}}, "('vae', 'vae', 'quant_conv', 'kernel', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "quant_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 8, 8]}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'count', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('vae', 'vae', 'to_nnx__rngs', 'default', 'key', 'value')": {"key_metadata": [{"key": "vae", "key_type": 2}, {"key": "vae", "key_type": 2}, {"key": "to_nnx__rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/models/recovery100/0/default/_sharding b/models/recovery100/0/default/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..b8e1e326dbe6145d0667351ef009a0b2b88d02e3 --- /dev/null +++ b/models/recovery100/0/default/_sharding @@ -0,0 +1 @@ +{"Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","Y3VycmVudF9sYXRlbnRfcGF0Y2hpZnkucGF0Y2hpZnlfY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLm1scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9tYXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLmFjdGlvbl9taW4=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","YWN0aW9uX2VtYmVkLnBhc3NpdmVfZW1iZWQuZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","b3V0cHV0X3Byb2oubGF5ZXJzLjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bG9zc19zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bmV4dF9mcmFtZV9pbmRleC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","bnVtX2NhY2hlZF9mcmFtZXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dGltZV9zdGVwX21scC5sYXllcnMuMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjAuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjEwLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3MuY291bnQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5kcm9wb3V0LnJuZ3Mua2V5LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLm1scC5mYzMua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLmxuMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5rZXlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjExLnNlbGZfYXR0bi5xdWVyeV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjIuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjMuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjQuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjUuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjYuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjcuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjgubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjguc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubG4yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmRyb3BvdXQucm5ncy5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkubWxwLmZjMy5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [\"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dHJhbnNmb3JtZXIubGF5ZXJzLjkuc2VsZl9hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, \"model\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5jb3VudC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS50b19ubnhfX3JuZ3MuZGVmYXVsdC5rZXkudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18wLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18xLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18yLnVwc2FtcGxlcnNfMC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5jb252X3Nob3J0Y3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMC5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMS5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252MS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5jb252Mi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5kZWNvZGVyLnVwX2Jsb2Nrc18zLnJlc25ldHNfMi5ub3JtMi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAuZ3JvdXBfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucHJval9hdHRuLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5hdHRlbnRpb25zXzAudmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzAubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEuY29udjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTEuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLm1pZF9ibG9jay5yZXNuZXRzXzEubm9ybTIuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfaW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmNvbnZfbm9ybV9vdXQuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzAucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzEucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIuZG93bnNhbXBsZXJzXzAuY29udi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18wLmNvbnZfc2hvcnRjdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzIucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18wLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLm5vcm0yLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5lbmNvZGVyLmRvd25fYmxvY2tzXzMucmVzbmV0c18xLmNvbnYyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5wb3N0X3F1YW50X2NvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dmFlLnZhZS5xdWFudF9jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"axis_types\": [\"AxisType.Explicit\", \"AxisType.Explicit\"], \"partition_spec\": [null, null, null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"} \ No newline at end of file diff --git a/models/recovery100/0/default/array_metadatas/process_0 b/models/recovery100/0/default/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..05d2ddfbcc01eafdd691217fe4c0cda746b266a9 --- /dev/null +++ b/models/recovery100/0/default/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "action_embed.action_max", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.action_min", "write_shape": [6], "chunk_shape": [6], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.0.kernel.value", "write_shape": [126, 3072], "chunk_shape": [126, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "action_embed.passive_embed.embedding.value", "write_shape": [1, 768], "chunk_shape": [1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "current_latent_patchify.patchify_conv.kernel.value", "write_shape": [2, 2, 4, 768], "chunk_shape": [2, 2, 4, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "loss_scale.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "next_frame_index.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "num_cached_frames.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.0.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.bias.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "output_proj.layers.2.kernel.value", "write_shape": [3072, 16], "chunk_shape": [3072, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.0.kernel.value", "write_shape": [42, 3072], "chunk_shape": [42, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "time_step_mlp.layers.2.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.0.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.1.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.2.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.3.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.4.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.5.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.6.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.7.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.8.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.9.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.10.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln1.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.ln2.scale.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.dropout.rngs.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc1.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.bias.value", "write_shape": [3072], "chunk_shape": [3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc2.kernel.value", "write_shape": [768, 3072], "chunk_shape": [768, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.mlp.fc3.kernel.value", "write_shape": [3072, 768], "chunk_shape": [3072, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.key_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.out.kernel.value", "write_shape": [12, 64, 768], "chunk_shape": [12, 64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.query_ln.scale.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.bias.value", "write_shape": [12, 64], "chunk_shape": [12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "transformer.layers.11.self_attn.value.kernel.value", "write_shape": [768, 12, 64], "chunk_shape": [768, 12, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_in.kernel.value", "write_shape": [3, 3, 4, 512], "chunk_shape": [3, 3, 4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_norm_out.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.bias.value", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.conv_out.kernel.value", "write_shape": [3, 3, 128, 3], "chunk_shape": [3, 3, 128, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_0.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.resnets_2.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_1.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 256], "chunk_shape": [3, 3, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 512, 256], "chunk_shape": [1, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.resnets_2.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_2.upsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 128], "chunk_shape": [3, 3, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 128], "chunk_shape": [1, 1, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.decoder.up_blocks_3.resnets_2.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_in.kernel.value", "write_shape": [3, 3, 3, 128], "chunk_shape": [3, 3, 3, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_norm_out.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.conv_out.kernel.value", "write_shape": [3, 3, 512, 8], "chunk_shape": [3, 3, 512, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_0.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_0.resnets_1.norm2.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm1.scale.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_0.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_1.resnets_1.norm2.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.downsamplers_0.conv.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.conv_shortcut.kernel.value", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_2.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.down_blocks_3.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.group_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.key.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.proj_attn.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.query.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.attentions_0.value.kernel.value", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_0.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv1.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.conv2.kernel.value", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm1.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.encoder.mid_block.resnets_1.norm2.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.post_quant_conv.kernel.value", "write_shape": [1, 1, 4, 4], "chunk_shape": [1, 1, 4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.quant_conv.kernel.value", "write_shape": [1, 1, 8, 8], "chunk_shape": [1, 1, 8, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.count.value", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "vae.vae.to_nnx__rngs.default.key.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}]} \ No newline at end of file diff --git a/models/recovery100/0/default/d/83878ea0eec4a9c18daeb44b066196fb b/models/recovery100/0/default/d/83878ea0eec4a9c18daeb44b066196fb new file mode 100644 index 0000000000000000000000000000000000000000..73c1499455421973db56bf338085a2725556b41c --- /dev/null +++ b/models/recovery100/0/default/d/83878ea0eec4a9c18daeb44b066196fb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cb3f3ae6412a95e2c6ca942b722c09a20a7f89c31f816dd0abca5234e87092ca +size 112810 diff --git a/models/recovery100/0/default/manifest.ocdbt b/models/recovery100/0/default/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..388f34cf829f4488ee4bdfef8d1366755c9476d1 Binary files /dev/null and b/models/recovery100/0/default/manifest.ocdbt differ diff --git a/models/recovery100/0/default/ocdbt.process_0/d/4d16686f37d7842d645df7bc0fbc9e78 b/models/recovery100/0/default/ocdbt.process_0/d/4d16686f37d7842d645df7bc0fbc9e78 new file mode 100644 index 0000000000000000000000000000000000000000..03e726a451c61203597ea7cd1ac75a63da892fea Binary files /dev/null and b/models/recovery100/0/default/ocdbt.process_0/d/4d16686f37d7842d645df7bc0fbc9e78 differ diff --git a/models/recovery100/0/default/ocdbt.process_0/d/4f2e215ac613b852f54f8856247bbadc b/models/recovery100/0/default/ocdbt.process_0/d/4f2e215ac613b852f54f8856247bbadc new file mode 100644 index 0000000000000000000000000000000000000000..8881173ed2e9d1912ef32d31828b5667f47c4c9c Binary files /dev/null and b/models/recovery100/0/default/ocdbt.process_0/d/4f2e215ac613b852f54f8856247bbadc differ diff --git a/models/recovery100/0/default/ocdbt.process_0/d/77cebf3655ec723496319cdb1c72edbc b/models/recovery100/0/default/ocdbt.process_0/d/77cebf3655ec723496319cdb1c72edbc new file mode 100644 index 0000000000000000000000000000000000000000..0007818f40d9ca859a5ba64c62cf62e272689ce9 --- /dev/null +++ b/models/recovery100/0/default/ocdbt.process_0/d/77cebf3655ec723496319cdb1c72edbc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:debe6d334fbfd88b24c3f1bf0eb8445c2d5ffdd63cc8f7f72a3a8536441189d4 +size 402022400 diff --git a/models/recovery100/0/default/ocdbt.process_0/d/989b03147ff1281329a5cc0adfa3cea3 b/models/recovery100/0/default/ocdbt.process_0/d/989b03147ff1281329a5cc0adfa3cea3 new file mode 100644 index 0000000000000000000000000000000000000000..e6af2962f5caba1e866287b1d7cc9482791e9bef Binary files /dev/null and b/models/recovery100/0/default/ocdbt.process_0/d/989b03147ff1281329a5cc0adfa3cea3 differ diff --git a/models/recovery100/0/default/ocdbt.process_0/d/d1816013dc8ccef6782787251339f26d b/models/recovery100/0/default/ocdbt.process_0/d/d1816013dc8ccef6782787251339f26d new file mode 100644 index 0000000000000000000000000000000000000000..a93767150447d4658c9d7f0f9fd92bb5796bbcce --- /dev/null +++ b/models/recovery100/0/default/ocdbt.process_0/d/d1816013dc8ccef6782787251339f26d @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7c6b80c42ae19b2a8a4fff6c8e85fca1200ba472d247c252316988106af97558 +size 177295360 diff --git a/models/recovery100/0/default/ocdbt.process_0/manifest.ocdbt b/models/recovery100/0/default/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..3f87d6f41fbe5daf876e4bb553d9a6dd9a74c057 Binary files /dev/null and b/models/recovery100/0/default/ocdbt.process_0/manifest.ocdbt differ diff --git a/models/recovery100/checkpoint_manifest.json b/models/recovery100/checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..b559815ba3602487df6cb2d4d1fd068d045fa9d4 --- /dev/null +++ b/models/recovery100/checkpoint_manifest.json @@ -0,0 +1,4247 @@ +{ + "conditioning": { + "action_dim": 6, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_type": "continuous", + "source": "action_delta", + "state_paths": { + "action_max": "action_embed/action_max", + "action_min": "action_embed/action_min" + } + }, + "flow": { + "endpoint_manifold_exact": false, + "inference_clip": 0.0, + "loss_space": "v", + "parameterization": "x", + "training_clip": 0.05, + "weighting": "uniform" + }, + "model": { + "config": { + "_target_": "future_sight.models.latent_diffusion_forcing_world_model.LatentDiffusionForcingWorldModel", + "action_dim": 6, + "action_hidden_dim_ratio": 4, + "action_mask_rate": 0.1, + "action_max": [ + 0.24, + 0.27, + 0.115, + 0.27, + 0.515, + 0.8 + ], + "action_min": [ + -0.195, + -0.29, + -0.155, + -0.365, + -0.39, + -0.81 + ], + "action_num_frequencies": 10, + "action_type": "continuous", + "aim_awr": { + "enabled": false + }, + "conditioning_source": "action_delta", + "conditioning_steps": null, + "decoder_head_mlp_ratio": 4, + "endpoint_manifold_exact": false, + "eval_min_noise_steps": 20, + "fm_clip": 0.05, + "image_shape": [ + 256, + 256 + ], + "inference_fm_clip": 0.0, + "inference_hedge_noise": 0.05, + "latent_patch_size": 2, + "loss": { + "loss_ema_alpha": 0.99 + }, + "loss_space": "v", + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "model_dim": 768, + "num_actions": 6, + "num_cameras": 1, + "num_inference_steps": 128, + "num_time_frequencies": 10, + "param_dtype": "bfloat16", + "parameterization": "x", + "passive_dynamics_mode": false, + "time_mlp_ratio": 4, + "transformer": { + "attention_kernel": "auto", + "dropout_rate": 0.0, + "gradient_checkpointing": false, + "mlp_ratio": 4, + "num_heads": 12, + "num_kv_heads": null, + "num_layers": 12, + "rope_spatial_dim": 36, + "rope_temporal_dim": 28, + "spatial_attention_per_temporal_attention": 1 + }, + "use_mlp_decoder_head": true, + "vae": { + "image_shape": [ + 256, + 256 + ], + "param_dtype": "bfloat16", + "scale_latents": true + }, + "vae_checkpoint_path": "", + "weighting": "uniform" + }, + "family": "ldf", + "scale": "b2", + "type": "latent_diffusion_forcing_world_model" + }, + "provenance": { + "checkpoint_step": 40000, + "converter_version": "training-state-ema-export-v1", + "launch_config_sha256": "1466ac6216c063b98404b0f5dc7c91afa426bd57fe92db1fb72734b3839a9c11", + "runtime_config_sha256": "1466ac6216c063b98404b0f5dc7c91afa426bd57fe92db1fb72734b3839a9c11", + "source_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_repository": "open-dreams-v2", + "source_runtime_transformer_layer_order": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11 + ], + "source_sidecar": { + "content": { + "resolved_config_sha256": "1466ac6216c063b98404b0f5dc7c91afa426bd57fe92db1fb72734b3839a9c11", + "schema_version": 1, + "training_contract_sha256": "f48d843e060cba4ec72237e7a5fff5cc406cbdac4ea14e63421c28eac0824c70" + }, + "sha256": "f1897d61274af463e7edd955ce53e10a2df5a6f6ffaae5554c0c9a111934200f" + } + }, + "rollout": { + "conditioning_steps": null, + "inference_hedge_noise": 0.05, + "match_commit_noise_to_label": true, + "max_cached_predecessors": 3, + "max_sequence_window_length": 4, + "num_inference_steps": 128 + }, + "schema_version": 2, + "state": { + "dropped_paths": [ + "training_state/modules/1", + "training_state/ema_params" + ], + "leaves": { + "action_embed/action_max": { + "dtype": "float32", + "sha256": "8aaf26769fba051bda5519a25115401bff1c6618ef6c8111a14ef487de53cf87", + "shape": [ + 6 + ] + }, + "action_embed/action_min": { + "dtype": "float32", + "sha256": "145cdab358e7550d07b3418e22e77875bc0dd93390cea924e3a9a11fe2e478d7", + "shape": [ + 6 + ] + }, + "action_embed/mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "927cd6920999e351caa3010114afe81bf3dbdf01fa29d846f44f77eeb4a30318", + "shape": [ + 3072 + ] + }, + "action_embed/mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "a3ed885d29d2466f7d8bc26cdfc54281567aeaf67671cc70ead9568cf46c329e", + "shape": [ + 126, + 3072 + ] + }, + "action_embed/mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "20d0bd7686da00dae302a14857f21b064071f009d4c02873fabf7f0508346ae4", + "shape": [ + 768 + ] + }, + "action_embed/mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "7f270c5d552bcc615030dda91b8149cd32e64110da854d4bb47e32eece617fb9", + "shape": [ + 3072, + 768 + ] + }, + "action_embed/passive_embed/embedding/value": { + "dtype": "float32", + "sha256": "6ae8f4898f45f5e002d60dae480a4c2e64622bad4d389ed1988d127e01116d6d", + "shape": [ + 1, + 768 + ] + }, + "current_latent_patchify/patchify_conv/bias/value": { + "dtype": "float32", + "sha256": "a2c5595859694297f415b6f85bf16ce885fb49e9019876825ffb4393ff5f55e5", + "shape": [ + 768 + ] + }, + "current_latent_patchify/patchify_conv/kernel/value": { + "dtype": "float32", + "sha256": "325dda5ac4a24943bd882fc80f50ef097aeb6fde76e3d16386f5eaedbd65f297", + "shape": [ + 2, + 2, + 4, + 768 + ] + }, + "loss_scale/value": { + "dtype": "float32", + "sha256": "8c472c1ddb95fdcef5759738bfaf0cadb6f6737488cec61a4fc5d10a9b8532c5", + "shape": [] + }, + "next_frame_index/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "num_cached_frames/value": { + "dtype": "int32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "output_proj/layers/0/bias/value": { + "dtype": "float32", + "sha256": "de743623f2e8abc27c44a360516dff22bbd73518910639908c371650ee46b98f", + "shape": [ + 3072 + ] + }, + "output_proj/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "797099b0f670a3190179f0b5228c1a8d21e185873d0c46444a3d398d206fcc7e", + "shape": [ + 768, + 3072 + ] + }, + "output_proj/layers/2/bias/value": { + "dtype": "float32", + "sha256": "c9d7e534edc4a516d530594f753c62dcd2eae78d5bc1468c5a6108713df157dc", + "shape": [ + 16 + ] + }, + "output_proj/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "586f0a1b99c93217606e1f0dddca10a88f97bc002f99890133451ba197402a81", + "shape": [ + 3072, + 16 + ] + }, + "time_step_mlp/layers/0/bias/value": { + "dtype": "float32", + "sha256": "66f4db5e5d00f335cdf77040b27a31b1bd794f63462ff12a4592b8a958953f69", + "shape": [ + 3072 + ] + }, + "time_step_mlp/layers/0/kernel/value": { + "dtype": "float32", + "sha256": "d87449f130a320f85316d448172efc2f1d8b96e9ead6be2d4fbf08f7482437ef", + "shape": [ + 42, + 3072 + ] + }, + "time_step_mlp/layers/2/bias/value": { + "dtype": "float32", + "sha256": "ac4a8ea885248b7499b6c49b5cc158569ffa8c6032c57fe75d6baa9ce935eaaf", + "shape": [ + 768 + ] + }, + "time_step_mlp/layers/2/kernel/value": { + "dtype": "float32", + "sha256": "eb9504510b819c870623602c2d415242bc51a65238b3d2f323e71713189a8f23", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/ln1/bias/value": { + "dtype": "float32", + "sha256": "a7d5b2995c2c8b5b07e76f18fccda2e0246fd6897a1e33bc3fdaa2c8687ba188", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln1/scale/value": { + "dtype": "float32", + "sha256": "a15bc808516d38fca689bd3f2735e6b8507e55164d585ee68fc52b4662736c03", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/bias/value": { + "dtype": "float32", + "sha256": "40def30ef077feaf89e163b9344f90214d532adedc98994653d18553a2facc7c", + "shape": [ + 768 + ] + }, + "transformer/layers/0/ln2/scale/value": { + "dtype": "float32", + "sha256": "1b3c2766778f45560f5c7749bbd7ff2e85476d4bde0eb234fa2b1ca38e50721a", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/0/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "49bfacd1af8fea5c5aa3e93f7b77b73b9b161410c7be528bdc41c19f9f66ab06", + "shape": [] + }, + "transformer/layers/0/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d3bf62ea4f6bbb3c20832063956cadeb6f7155d6fb8fbd6589b9bcb3cb4a0501", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b959090da8cd5e744cb1dae97a5f7d7c9e9ad437d284fef6939e44ddc0b62612", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "718642ddcff89231e25fca0d0de820ab42e81d9e16a2f1dd7f1bf8181ae0ddae", + "shape": [ + 3072 + ] + }, + "transformer/layers/0/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "24233b0fc3e7cfa3d3c29cee4a04cb6bca60aa4ab0c6477ac3902e7efbcaa605", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/0/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "487987571e49ffe12edbebe4105038df2ce44ebb9449d0d2b8bbbf1c92c89c42", + "shape": [ + 768 + ] + }, + "transformer/layers/0/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c67d52c8277864fbbfdd64f4a01dce1dbd097d080ce3205d31f3d3a33e001752", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/0/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "3f5b5ff8792b7d276e0e9473c1b693b329dc30bbe2f2264e89a4ae4ab66df3e4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d2b47eeaef2acf37558356f3408a691fcc213033f4cd7ba1f7e2df6ef74f88c8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ef47f45013c1c260e754e7837324f25e19c248d9b171d07ade90755cbffaa258", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "87a2ea62b95d5803947a2c6cc4631036c9e0655f09b8174b9b2e87eb10103c8d", + "shape": [ + 768 + ] + }, + "transformer/layers/0/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "b76b24e45ff106a41c0a2d31d4fbe45cab0136eb04335d09b09f850ca559194e", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/0/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "e6682c39234eb38b990294e9169ad12e70b626f2031dd257f4b932599ac8588f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "ef0d674415958e46516dd731fdbb6e33620e77580490a0faab07adbfff759538", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "cdb404e3ee0ac2355cd1f2c84ab4e81f547a1a369822bf4f44c4ce7981fc0841", + "shape": [ + 64 + ] + }, + "transformer/layers/0/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "40146abd460ef1a512955bda53871c531f9aabafa56fecc21b882a1cb2e9eab1", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/0/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "64d7d5ce3d2ce93f5d916feb2f1b33d269a681a68a643bedd012a9358d9d5e0d", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/ln1/bias/value": { + "dtype": "float32", + "sha256": "b177ac5d425a6c0ddaf1b2eea927f920cff86d584e76891c2fbddd9a1eaba117", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln1/scale/value": { + "dtype": "float32", + "sha256": "463e4c58c8997138eb75bb49f746c9b5c8dd168970803e553063c165627046a8", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/bias/value": { + "dtype": "float32", + "sha256": "a8041917eaf8fdf4c0c38a38a4d50058490207dc4984546fc18f84bfe72d6e35", + "shape": [ + 768 + ] + }, + "transformer/layers/1/ln2/scale/value": { + "dtype": "float32", + "sha256": "f4e3bc9ba16073c1df26adcd0b96ad52cccdf977312f934bc45c092337aca1f0", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/1/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1b6ee3f31d4d14acd6eb5ad151fee556456b96bceb100476913057f48cc526f2", + "shape": [] + }, + "transformer/layers/1/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "88f58bf80331390fa8119fc8eff7d1d28af5b6a0728249ab47bee8cc8087e380", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "3dc320fe074cfbca08b630e10fa1caf3ffdfedf26069bfa75d877d0d8569dea5", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "b4a47c7c9397f00358ac3d208a3083b8df6ff9a68777b4c736f1e60b88403939", + "shape": [ + 3072 + ] + }, + "transformer/layers/1/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "882724edb1f52e31e22ab3eb9fca5d3ba6c4551b67c15d4293ea52f790d9924a", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/1/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "c431f6aed2532eb0025d2c51fa3a520bf8cd1589f47a02f6247fba785b0f0828", + "shape": [ + 768 + ] + }, + "transformer/layers/1/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "e783d50067b5bf2640772605cf171c8e5c2c7602d0e870ec238979b545ed959f", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/1/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "14928c8696bdc17c4d693086844ff66192c46d8ccfb50d90fea5daaa6cf3cd28", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "9b5a54b5eedcf777c965a1302be6f5335317b13a0875157e998eea1aa35dfa50", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "157d4f539956608440b7cba8c5008366b9bfe014eb3e7ae70f0fd8098364d7a5", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "c6625cb2685295b8ea91342dc4ff25909c1d90d29101c283792275051d584333", + "shape": [ + 768 + ] + }, + "transformer/layers/1/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "2a49f6d0355fccfc92160a86ba6b2312aacd2e3cae8d0766be93a517712fbcaa", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/1/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "8726a7f5243794f9193fe7a96c369e518c11217083bb8559c3ad38ad86d8eadf", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "dc14c78c455e6381a12fce2845f153aa1402de0bd530bde08a4aa3b63ada437b", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "944d028d7128a429db666b2362e02df0d6c448bb6063409aa638ce846747d63f", + "shape": [ + 64 + ] + }, + "transformer/layers/1/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "e7dc07cdc5aa2b6676a2836887194575e6d4e534cf8dcd219d236d297d06342e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/1/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "91ad8ace27cee9c5aea8bb0e46d7f0511f14e5e64077623a0a054340eb901592", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/ln1/bias/value": { + "dtype": "float32", + "sha256": "1b8bbbb0e6801dd78e5282ba113ff526a8c36fa70a780f6a5e06ce45804103f9", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln1/scale/value": { + "dtype": "float32", + "sha256": "641be821e0afe87f2e7b553fcea3267da881f484ca3ee50d82027141e2abc075", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/bias/value": { + "dtype": "float32", + "sha256": "92e5fd1e497799c291717933ae3ffc7c3a0feafcb11df94456028311f0678a68", + "shape": [ + 768 + ] + }, + "transformer/layers/10/ln2/scale/value": { + "dtype": "float32", + "sha256": "64b83cfeab89d42744c8c10990a102337661f2b7ef66d53b4a2e7543ef7f956c", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/10/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "0bd5c5a3d09a6888001516b03c63901f1cd71763cdaafe6403d4f92bf21679e6", + "shape": [] + }, + "transformer/layers/10/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "81e24bc6fc5e2620417c379ee0d5f5f1e1bc9e920d01ce8278aa766379da7c9d", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "857f289682a5f0584b2d1a3860fede0c6f1a19b4dbcc6af6ec9f616bb281cfdc", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "852b887a99c13348ffcfa9fc945d5309dec235a9dc45eefc483f90ab43fd964c", + "shape": [ + 3072 + ] + }, + "transformer/layers/10/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "b5a3405d475a9fb81e103d48173b7cb13ac330215374c9a08c92515523e065ff", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/10/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "711497c88a2fcf150588d6f0e1bad0ae5c769b0382a4d5a6f6ad0b995a16eae9", + "shape": [ + 768 + ] + }, + "transformer/layers/10/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "0ec8717273484b09211d6c014e51bd00f63bce54e65e1a58ea9935f75d989a17", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/10/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "ded62d43dcd865e37a0f5681e300402a1379aa0b0edab6df0d570d42285cc34d", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "89982809806050e95af56775cdab706dd341ee34b7de3d00c1d20d5700ebff0f", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "4f4918360eb72d8ee052ea485528d92a6585dfb489be7bf7f25414b146aac4f6", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "be5121779e879b85ceb2827b59d9e94a3656fe63b9d218b7f3f293f13ffe34ec", + "shape": [ + 768 + ] + }, + "transformer/layers/10/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "869ff31ea6088a8fbe9a532ead74d7b4e89f9536ca57b9f327a5c6e2c96eee26", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/10/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "625f37c849a7dd6acaa5363818a246e6773933081f1bd1cdbf713f17660977d6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d76ee37118457a7b85d01a26a31f9d6da083705ec52bb6a1808dfb02e03b9103", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "76aca2578b3de09fd7ba4774e32647e73dc602c10302e8a1220a9f1c39cce26c", + "shape": [ + 64 + ] + }, + "transformer/layers/10/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "a43c91d2eeee61a857a83750e29093bae24efda023466e1eafa1250c85ab0e12", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/10/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "53788ff32fe4b6ac0bc470a8f8244d1a61a92dc26ed9ce3ded985f507011a0d8", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/ln1/bias/value": { + "dtype": "float32", + "sha256": "5729f073fbea062d42153c6765c35dc01dde75b952b56a8102d09fd879fc7d98", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln1/scale/value": { + "dtype": "float32", + "sha256": "c965fc68defb68fe6ba4b0a2bf381d56c56f5e46676285a12d4b5eefc459fa51", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/bias/value": { + "dtype": "float32", + "sha256": "3959c8cce8d38b89c69b538fb790adbeba8813dda33f57c7d52cb76cf6b88e7a", + "shape": [ + 768 + ] + }, + "transformer/layers/11/ln2/scale/value": { + "dtype": "float32", + "sha256": "0813516fc2204e647c0dbfaba09c01cef7bf12f711d3164ed26b6eb8fc6c68fa", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/11/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3474684072d2d9d3db0eba63d2ab92a70e70c259173f667f193a48932d5b0dda", + "shape": [] + }, + "transformer/layers/11/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "cfa1e6f3b5073c426418f7f45cbcd258ea2d81002c94862b5eb24740aff2555e", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "cee96144528a6f114c6fecb5c65b8bb3b5072ae97fc7b73864aaddc162ae6283", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "58aca6fbe0f472711ed4858e315a7c177f72fa286fa69277683808a411ae5742", + "shape": [ + 3072 + ] + }, + "transformer/layers/11/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "92ebe21168d97cfeceb9d1fe51bdc7c07114f4e03239decaf8714cef6b5e891d", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/11/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "a5a93fc11c150ef8499ee5a77b26ee0bea48848d7ffbd09dd61c55198a9c2f69", + "shape": [ + 768 + ] + }, + "transformer/layers/11/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "6b806b22b57b6209fd3b74f55a54a346fd9c7a322804ee838e8c89562d236288", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/11/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "512bcbf252a79bdb19b92091c52f8a054971186fd6818e95f4d9a2921b7ea55c", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "bfa06244ba1c5b998e70c0e3f3e27ab623d4ff174220ad0265034c6fb3f6e14e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "2f8ffde5cba929ed4282f69d1d53eb3abc6a025d76abd087615eca9f5b18c02c", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "47f222b776ba72deb2d4634a95ad0d0a5a51c34a3cb9b94a46db0002780bea25", + "shape": [ + 768 + ] + }, + "transformer/layers/11/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "fda28c20fd7db98bd084a675af91c8ca948acd6be81a9e1fb1c32dfc1a0035bf", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/11/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "0ac909261e1c105fe8afa7892ffc8cc83e0600d2affa2a3bcb248df5230aeeb3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "3fc73bf24086298f8769fcd13990a401690b0c3c413ccdddd440ceb6cbbf26fa", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "631f83d6f6a37be692372b490d25fe1e7eacfe7e39793ab33985087722e09371", + "shape": [ + 64 + ] + }, + "transformer/layers/11/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ae275971ac395677c44bf07d8c4b67d5121bba2b93468734e2e4c3731f4fbe4f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/11/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "4e65f34a59215893606cc6b034fd1643e168003f106a19c0ee44c5892d3344ba", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/ln1/bias/value": { + "dtype": "float32", + "sha256": "eb1223ef84f52a040bff8215237dcb73b538bc684f0dd90a0d739ab944a290c5", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln1/scale/value": { + "dtype": "float32", + "sha256": "44d4a94d2f6b2268cfee5a56e41159cd0baf2f03aae1a1ff9612894313b1d9c9", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/bias/value": { + "dtype": "float32", + "sha256": "d836f54fcdcc0554afb997f1fa1e94b5e7724a905fc06b7388557572a71ca29c", + "shape": [ + 768 + ] + }, + "transformer/layers/2/ln2/scale/value": { + "dtype": "float32", + "sha256": "8103ed523f64ce841aad9fa9d982d8ba6053a0097add935cd966cb8257940f77", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/2/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "1cf24543199f7c9ce1612d44d1b51dc46f8547cea4b3fbe0aed41e89277e9bee", + "shape": [] + }, + "transformer/layers/2/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f10cedd7f8a3670a7edb1c1af97f3b5c9da895b2428c6e34e7cda7efd9acb102", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "34d73f6f76469035325ab34581ea5cd5765cfdedd797a39e64ecb571c88d747c", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "a34dd74a734a30f7d07afa56714434b0959c567d3292c1e4f3dd7f1b285d2450", + "shape": [ + 3072 + ] + }, + "transformer/layers/2/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "f9905f2a6109f7ec292f35f12b462e311f6c8896833ecb1281276e696b46a94f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/2/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "4d06ba80d3ebe98c15b87c50e1a62d59d8a9254f4d8dc2d732e0fe1ba32cd36d", + "shape": [ + 768 + ] + }, + "transformer/layers/2/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c68353901a22defb22d188e89312495c387e62014986989e2adec437c1a0c48c", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/2/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "cdb53bf7840b6066fbd1f479b123ae3fc1b977bb969d2c5dae9172c37cbddcf8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "7187420aaaaa27256a1e5c7098f2d6f18b1dda10f12e698d8dad4ae57981ba1c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "733bab310652c0a439eedce05098ed77a2f2ac7571886524c486e0fe0f1cbc4b", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "4be39ac241e5253d40ddbf3e1cd775eab4fc328da4dde55f61a14f6c21d764f3", + "shape": [ + 768 + ] + }, + "transformer/layers/2/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "ca2824d4ff23c6b74c66612ddb1f997b982f36dc0b3462d65c512e7d7d6fd269", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/2/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "bbcd4acd187d90267bc4797005829a936bb5e463c3c2d658fb88945ffdf69815", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "e6e1c8dce0d694304b6cdefa6de57a28ed09b5a48a27c663bcc2c6650e2c70c0", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "100518363f79707051157eb3877bac1db94e4b0ede17f72a43e3e2ae440080db", + "shape": [ + 64 + ] + }, + "transformer/layers/2/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "d46265f7556aee5d1a2b055ff1b9b7f7bdb318cf95251a04447ab7ebcebe83c9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/2/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "2cd8b470538c119a5bbda5da57a55e4afb94ec3c6a7e4e4a8c2c39e4a5deb2d3", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/ln1/bias/value": { + "dtype": "float32", + "sha256": "26f1860cffe8a674edd0603f76e14452291e58bdca92660e41a49702fd66f76a", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln1/scale/value": { + "dtype": "float32", + "sha256": "182ef342b5046e9e9a005cca8dd708cf7f2972ce593148cbf2bc846f80f6b1d0", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/bias/value": { + "dtype": "float32", + "sha256": "f9d185ef8775015617a931f68efa100ba2f96659e0337d4bff3fac1fcc43b598", + "shape": [ + 768 + ] + }, + "transformer/layers/3/ln2/scale/value": { + "dtype": "float32", + "sha256": "8a4251d679aacc03b03562ec24312c64e0ef532acc0f7bf96e825b07aa6c1d5a", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/3/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "57c871ee0899aa9bb1de5cce329dd26ebdc5da9292a967935960a05ee9e7d638", + "shape": [] + }, + "transformer/layers/3/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "843b7b840530426cf77e25d12802a4005458b6b4542b00ee6b62842de716deb5", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "7dee0d02c88298bd3b67c0ff042d482f96a7823fb34f8abd7c43cd1c351c2237", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "5e77d6bd2bb8d1ac1308f163ec18f0791cea14695f18429d734bd5013f5f1148", + "shape": [ + 3072 + ] + }, + "transformer/layers/3/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "7e6c2dacb03b86af6c3aa4eadbc70f22c1a1ce9a1e3e09214525839a0366b242", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/3/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "73f63e2c6032ef7a32def04bc8aa10e155001370ee4ff47f8e79d2d9c82cc4bf", + "shape": [ + 768 + ] + }, + "transformer/layers/3/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "1920dd51b1910fde21c0c81cbfb22214d42ef88fa2407ed375fd4f994c549ab4", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/3/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "e73b383076de1c8b38ce41f601b16450fd731616f0bec242948e1adb27bb44d7", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "2edb1e2ce295a680f1b9f4eb968b6351e63f8df3fb49276b254d9d6c0d485bf3", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "883ef3621c65fee98cafe8e8b29d9c930a62b38b2ee9ed269c2a5206e809c18a", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "dad1f224598467dbae49ce6d9145c0d51dcd729ff8e8cc06bd66bcc376cb1171", + "shape": [ + 768 + ] + }, + "transformer/layers/3/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "09c472bbcba7cbf8c5fb9aa6f2937084ea8bf3ca00c593d9baa897bba98f2b80", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/3/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "037f58af52f634fc057c04a52dd463b8ce7060aa4831710963907fb6e5754090", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "5b2372596dc0b0d98dfd575d23d97755653fc87afd0a877237f4831d36234e4c", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "d536386ddcb35d6f4a0706d5e697f2929f57a76872c8cbb37fed1e11a90e809a", + "shape": [ + 64 + ] + }, + "transformer/layers/3/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "a27aac14e9bdee4b75876f9bdb55b18e31e253d9ade9549957b8b54e1de85fb3", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/3/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "67cd79b670f14012c5e8beda829a1af40cb62a0eeb66aaf1cf19b9dea390f578", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/ln1/bias/value": { + "dtype": "float32", + "sha256": "b0d4055041e72ee515f6f2380c9414f3e82b9005c16b0172c48e3fc5f2079c26", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln1/scale/value": { + "dtype": "float32", + "sha256": "eb059da7c6998eaf5593c822de55c1948342225b0dc7f21d3d8f9233ddf73fe6", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/bias/value": { + "dtype": "float32", + "sha256": "4b172a8ee529dd7bc4e0296675dcaa1c85258628a564e8a86977fd20669fa065", + "shape": [ + 768 + ] + }, + "transformer/layers/4/ln2/scale/value": { + "dtype": "float32", + "sha256": "f60b3ed9d4ba5d511bc5119f4978af3c16c3eba0d1657da694879b585ff87670", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/4/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "9673d6e11c9e6ac6935e19032e73c783e868ddc1a04a93819bc97cff6af2740f", + "shape": [] + }, + "transformer/layers/4/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "48f01b1153b529497063bfa3c4382c1d728c33687e8a0bb3dd7593a3d193bcb4", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "b9c52ebb896f43f852b9f3656411947e200e584e53ad0b204ba297a6fa886228", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "774039bf0330187ff0e0a9063fedd5f6935da454b4935121992bfb8d7d4f8165", + "shape": [ + 3072 + ] + }, + "transformer/layers/4/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "3835d739f1ea44ccffb05525e362c2e619210cf6bf671084cedfab26dddc5ac0", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/4/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "5bb8e1c2bf0e1788788989ad1aaf98314ba4095f936399dd2de423a1e7c8afcd", + "shape": [ + 768 + ] + }, + "transformer/layers/4/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "ef514a0d249f821a25bd112c481ff095ef611162b857fe80381c9d0b906841ed", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/4/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "aa938eb727287e4c94a99c83abcd7f51f1e9f0a48e8805214fe50e8f0220ebc8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "43ffad15f8791bd30730d321a5c58254d0b869ee7d5aff113ba73066ed505e96", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "fdf2402f0406ce04decaf5bc810be5e9a38d491d153bbde9598dae387cbe7f3d", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "10df1d999df3442f18cb86a133f3dd2e957b5e5bbb0e535425d4ae678bc9c296", + "shape": [ + 768 + ] + }, + "transformer/layers/4/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "594d2dc3fda3c090e906f36f518c6cfa15cbe9547d3feebacb64003dac497bff", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/4/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "5b3e54fa021db0e756e2eed21bf4ee40218e82cdca3754f1f2891b693fc485aa", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "d2e0f1e127b7267642017374b86b96bc941da33654749d7779bccb6b7ff9cec2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "4353f3251f5108809a2a86887067d9a5771cea27dab62930165748c341b56d48", + "shape": [ + 64 + ] + }, + "transformer/layers/4/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "59432ce92f1c70e917830be2716857eef352e98d7eae67ae0c876eda91b1ad8a", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/4/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "da141e8c31cae462bc6f0c9063c1e6e5f6d319ed53b2fb859dce2e1e56c8c8af", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/ln1/bias/value": { + "dtype": "float32", + "sha256": "756447b7f10fe098fbb4695a7c60fee9d8c9e1be4c548775f917f3a269a56b56", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln1/scale/value": { + "dtype": "float32", + "sha256": "0de9fc494b71b75533209c835841d143596539c189dcfaa4b8d4fd823ee205f0", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/bias/value": { + "dtype": "float32", + "sha256": "c74623963574554a41ecfbb792d4ccb2b04dfca779dd5e9fc0aa6ec7d84da841", + "shape": [ + 768 + ] + }, + "transformer/layers/5/ln2/scale/value": { + "dtype": "float32", + "sha256": "ed2e75ba32ee1ef40f8f54b323c1d658730b07eb71f40f645cd1561b1587a7d8", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/5/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "3599f1941923f09c21c3e397c473a8706fbca94842e9395a4f61e4b18149d41d", + "shape": [] + }, + "transformer/layers/5/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "ba39cd0eead86ea77c1d24ed0e03783a56adb76b28ef71665740322e66cfd0e0", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "f75dfa06eabb4d76d148b0d72ba2707c4f3ab4cc9fcb2130b61f1107b580000e", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "3b2a68b3e77282a4fba402c148371b46ad030b980537bc7c4b3621d7c6b055d5", + "shape": [ + 3072 + ] + }, + "transformer/layers/5/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "c2823eba231d70c0ad3ff19cea433da423efb5c120d66601b159e9dd495cb850", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/5/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "3c8a81bb0c6ad0ddbf2efb47ec6a1e8288e0c1325a668b99d49106a078918407", + "shape": [ + 768 + ] + }, + "transformer/layers/5/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "20f8c176af7873ece31187d3de08c2996fae15ead7ae15d0bb44e95a18cb4c93", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/5/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "f4afca2cd056fd85ceb4b0f369f2bfc93bb3e4f71f157d3d5c31db6eeea933f6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "be0d43b3e9efebffb7fc9b6d32105f7a69d53df563e763b4d566130d36de62d2", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "a80c011ed7d819361a2e7e4ceef225f2880952b57f5ccfd92f892f5cd3191488", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e448b900dd518c3cb8696441b1085e2c32bb4824ecd26f8736dcaa67c6343b2d", + "shape": [ + 768 + ] + }, + "transformer/layers/5/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "b7d0d48ad621c25d5d2a353ce9e751a5a7540f76a37f3283b5d856c9a54691c1", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/5/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "f5e53ffe3e8d4be16578a19e1224e8e8f494aa25971fb15c559b07240efcae23", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "22e8e23a46ac1f9a960516a86a12681fcc2e14cc9b6d49964b328a74325664ac", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "dad69fb01dd238835c4e782c579a9aebbdca688cb05dedc9902f1e707e908e21", + "shape": [ + 64 + ] + }, + "transformer/layers/5/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "79a54db0c59f83aa6e23c2e3194ae45ec144f8670cc25ce3be098da7d1c38c68", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/5/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "5e99fd1d742e649101cbdfaa5c49efe11d5407c06232d406519cb6572495bc48", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/ln1/bias/value": { + "dtype": "float32", + "sha256": "5a3029e5e4fb5f6107889112453674160936a9b2b5f7fff29cc012d805514983", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln1/scale/value": { + "dtype": "float32", + "sha256": "be94840dab3f12c27f339f151da9507571a9b83e629a1e2c4844f1d9cf1b57cb", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/bias/value": { + "dtype": "float32", + "sha256": "57728230f3444a96e0167487fe60763e7245f49d6b0684f4e5c38e0523581357", + "shape": [ + 768 + ] + }, + "transformer/layers/6/ln2/scale/value": { + "dtype": "float32", + "sha256": "e383016d82af153cee55c6f581a69712a4dc0474f97921e05728ff82de677b55", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/6/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "f600cac463e977ead52b0abbc6f0b3ef3c53c3ce864291282b18b5c66b2da3d3", + "shape": [] + }, + "transformer/layers/6/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "f6d59a832f3beacb2d3a770b7d5331888854d80788c17cbbd6161b1c75fac0d5", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "e378670ca6e9c8df5fcb655bb35ea943b7cfa5006e21da7a567217a225d7b5e6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "6dadfba815fe74871fdc3716c1e2b68ccc2ba3440017b3712f72cb037efa1a8b", + "shape": [ + 3072 + ] + }, + "transformer/layers/6/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "dda058762f265c173d4d39b2a2cd76abf854799674efeed4ed5e514ff8e29dce", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/6/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "26707b56c79598e9fb0eccc9ebaf0c1974158832df501a48e61a46a4bd48c295", + "shape": [ + 768 + ] + }, + "transformer/layers/6/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "9c4a2b2552cb47dda57481a489b64b3e460d7bdd3926db5118d4d32498a7bdc2", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/6/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "7a96d4df49065fe1175d69f06e95acfa0cc422899316dcf4c2eb6ac729a963c6", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "d393a837d70562ebcd289cee01a860ad9036e86638df15b5768fd63aa2fd7f4a", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "6d7b157bd844a0f3fb7bf9829ed9a337a5938526df82659d417261b2cc92857a", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "09e30c1527463db8451fe439967ab1c104e11faacac7d19e9df8864b0833a885", + "shape": [ + 768 + ] + }, + "transformer/layers/6/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "e5a173cb991e0180ac8bb77bc71a58751a55b3b4ce1d2427f08640f5176cf5bb", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/6/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "d7ceb84c7722e311eb6756a940e6e8d912c2dc8cb6e6046121fa6dd94e39a51e", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "25fc1b2431707ea68cbd04585183291cd926e2956265ed2809e4bb006799338e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "659f4657a791e02fadc4b61d203401a9085493da5dea2a9f60aef8f494bc2103", + "shape": [ + 64 + ] + }, + "transformer/layers/6/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "808dc9938c13ae1c8f6bef5902e5371ab9fc9fb54ab01b31b6684da35a7edae8", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/6/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "060c106187685f3984c2ee0e864e8b4fcd4780b27ff872c2d374d8f9b5e78f01", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/ln1/bias/value": { + "dtype": "float32", + "sha256": "be544117880232a06fde0b8523640295ca365eb44fa705793ba2bd1d5a53223d", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln1/scale/value": { + "dtype": "float32", + "sha256": "9a0fc08a0c46c96c2cbb30bc0aec888954ec6c8773c17d09f67ac23026e3b033", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/bias/value": { + "dtype": "float32", + "sha256": "bb17614ce3487944eae85caa8dd4f1d1d5bcc51966d8bce6302cf081f8ba01f4", + "shape": [ + 768 + ] + }, + "transformer/layers/7/ln2/scale/value": { + "dtype": "float32", + "sha256": "00334afc8ae57da9d4203e76a58fb40fad4303b93430ca281c295fcff28280a8", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/7/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "96faf1f6b10b22419c64af3cdc125b25e80c0721ad6c51a6325d696c3c1f4a1e", + "shape": [] + }, + "transformer/layers/7/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "e8ecb9db26bf9c77049e9f6cb72517eba24ca57cdd378307182aa66be0492107", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "89315954c29be19257167ad1a8063297a5461484a0b9a920c1491ed7c1780f96", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "d71c6e0c5c98cdc14860e2bbf4866e43737dcbffca02939d4c71ec5a639b92fd", + "shape": [ + 3072 + ] + }, + "transformer/layers/7/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "071711a307c39ceef05f2990828d4730f8658ed69e36d5955490f512bf2b204f", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/7/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "89c26efc0d22d9832c9125147c413dd7d9e20124c195f955965bc89d82a50c5b", + "shape": [ + 768 + ] + }, + "transformer/layers/7/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "ebac748ffd01ceb781a04e9a6a50b54fbbce4fc2eb4bd81b9c9afed34d2b3bcb", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/7/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "cbbe83b7369de9b18a14c902815b7e8f3b9ea86f11acbb0496ad871f07a2d3f4", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "034d80c533196118f8e7f6e24fa96d243d0f909be1e6037bfbf2d29db8c91814", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "f232a49c212a0c170a23e4725c9c05d7d9996d8581667388cc62998bf149cfdf", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "e23781a41663e556ed1fb528b0248470a388be01cdfd859b1da7e0c02e697723", + "shape": [ + 768 + ] + }, + "transformer/layers/7/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "fdf7a145244065eae8f4402fbe19949162b00ead37c312bd5a5a77918769af02", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/7/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "7951e706390523e227d5ae7ca605024a416bb7536cc76a1485427ec08240dd7b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "3b8ec4a6a6a6e6b82efd038a3d027906326cb1660243f5d9235ce6e5a1bfd0c5", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "3dae5d121f0df08fe1c615d207b897ab56d0561fc7a53558d9c8b421906c79ee", + "shape": [ + 64 + ] + }, + "transformer/layers/7/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "92cc0656487db67ccbee24f4a18c3f616e27a3ecca99f1410cc2c4db4a78ea2b", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/7/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "0425e5bae25dd2639bdcec9767987bb19f9eccbde9c908f6595748e0712f8aab", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/ln1/bias/value": { + "dtype": "float32", + "sha256": "6cb6ab94b6ec0f51418d0d992935f51befdc67370a3ccbf8441d572496cb3bcb", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln1/scale/value": { + "dtype": "float32", + "sha256": "c04db6e183b7b2ab1182a58f266865a5518c6eb87a4f7cffda378d06505e96fc", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/bias/value": { + "dtype": "float32", + "sha256": "09b11911e71f88845f6df5fd1fab48f95bb89632ef98e937ca612b543cd8ed6b", + "shape": [ + 768 + ] + }, + "transformer/layers/8/ln2/scale/value": { + "dtype": "float32", + "sha256": "3e188a426bee29409346cf1f804ac7fdd25d13e7b877afefe98c5636acfb3f22", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/8/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "a8a11533153c3c4467399686f44ef6e1ba2994423c677d4dbef12ec15085d7a7", + "shape": [] + }, + "transformer/layers/8/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "be978f8d4ee3c100cbf5e0967c8613e67f6684e838a817b33655c638580e3ebf", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "c93179b0426387af7d4b62f8c0538582424dfa90dae8b9b01b87b0a5ce07c327", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "f992f04e3464e578173ae63e087cfd001836223d34c09d2a29c882b69e22923b", + "shape": [ + 3072 + ] + }, + "transformer/layers/8/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "119ebadeb2721557a7c4858bdcdf89763161ea5ce724e44c3fd92c2a4395ecc2", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/8/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "e61387e13956838e72fecd4758a2714dd5ad5eee8b88dcc8daf2a6576e5106c6", + "shape": [ + 768 + ] + }, + "transformer/layers/8/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "a5fa4ecccd66893c82e3a5002f1f3104721767c32af0fc29aadf3413fdc9774e", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/8/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "9bd511b3473cbccd0ac020840dc39fddbcf9f43b3f4c79b4dc2a151b77267407", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "f5d25f8931dc03d495a43932ec87b41078666438dc870861bd560def8252e301", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "ef284b2bce605bf71a9ed9fd3eae6f2f7d7ab20f0a5ba1e85f74f8418a6cbc08", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "9b9aea947565091ff790258e05b45dbb905cfd7f6a19c0c9ceb824880676a528", + "shape": [ + 768 + ] + }, + "transformer/layers/8/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "bf96372a2f506c51de26c380d7a86fa1131f69346d782692dc4b8d8241f1ee28", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/8/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "9033c6ce6017fdb9e08cc185d165bc0125172f3700a6ae86befcfa538bc73c18", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "53389ec5666f0fa2be39b640be2135695f291606c91d0c1f63e72b4e1c36024e", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "f8df4fa59741855bc30be5f9ff056fe5ddf32e2d8638b081f7ee10cac7613cc4", + "shape": [ + 64 + ] + }, + "transformer/layers/8/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "ed18614caf1e781a2350945beba0439333353497e86b3510914fcf4179324aec", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/8/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "781093e5fd1010ccbd56e773c2d6b9be156cee3eb0f1b3630362757e84554f72", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/ln1/bias/value": { + "dtype": "float32", + "sha256": "c785269201ab8ada09b3edbff8095aaa69ca0e3bba3c44bb65231413b796b8b0", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln1/scale/value": { + "dtype": "float32", + "sha256": "407c69492424dcdde2b7ffbda01ac449c71f0f834a0f176181f0604bcc7a9cdc", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/bias/value": { + "dtype": "float32", + "sha256": "a9722ffeb54305c9876d92d4044a3f8355ef0f7c9d22691fc1ab4c2dbeb67a44", + "shape": [ + 768 + ] + }, + "transformer/layers/9/ln2/scale/value": { + "dtype": "float32", + "sha256": "d7f37caddde601cfb66e8707751bfd313f9d3372a20bb6b9db5d86c62822eb52", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/dropout/rngs/count/value": { + "dtype": "uint32", + "sha256": "df3f619804a92fdb4057192dc43dd748ea778adc52bc498ce80524c014b81119", + "shape": [] + }, + "transformer/layers/9/mlp/dropout/rngs/key/value": { + "dtype": "key", + "sha256": "671589331846ccce7b790600a7d511721f513372fd669879b957396b0914b239", + "shape": [] + }, + "transformer/layers/9/mlp/fc1/bias/value": { + "dtype": "float32", + "sha256": "d923143493901209f495d2ae802b006335aaba1c1432d9877225ec26565b230a", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc1/kernel/value": { + "dtype": "float32", + "sha256": "954c2391587cd2a5c463d2f88d409f19140d59a67e77100d32f0e8fe94ac46e6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/bias/value": { + "dtype": "float32", + "sha256": "e3adff7fde0a8c7635bc0efb1ae5a4fdd401c0a7d17c4104ee95eb6839981989", + "shape": [ + 3072 + ] + }, + "transformer/layers/9/mlp/fc2/kernel/value": { + "dtype": "float32", + "sha256": "944e9272a37dba537a3573bc99ab0181398765682559fba57bbac68a887c3fc6", + "shape": [ + 768, + 3072 + ] + }, + "transformer/layers/9/mlp/fc3/bias/value": { + "dtype": "float32", + "sha256": "ee9c20010dd6db0f19c67abca4f69ab78644637d33f44ed02aad39015dfb8a83", + "shape": [ + 768 + ] + }, + "transformer/layers/9/mlp/fc3/kernel/value": { + "dtype": "float32", + "sha256": "c43b59f72794b46597d6eb6a333d6a6aaa4fea81b96b0e7833b2811957b39e2e", + "shape": [ + 3072, + 768 + ] + }, + "transformer/layers/9/self_attn/key/bias/value": { + "dtype": "float32", + "sha256": "837cc68373fb44e467ff4e9b56a7da595e37392cd44c42fa161988d86b678029", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key/kernel/value": { + "dtype": "float32", + "sha256": "e41e32068236e5cf0d8d7b6e2e866139516d465dad65b7e2f48ac103f2529406", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/key_ln/scale/value": { + "dtype": "float32", + "sha256": "1544debb8f1dc2f9362f7878938795bfa8a9871e2c732e83a6ca15a254006099", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/out/bias/value": { + "dtype": "float32", + "sha256": "5223b174cda833542335017cc6b2fc8731a3f1964b1441bf14c39caafa91332d", + "shape": [ + 768 + ] + }, + "transformer/layers/9/self_attn/out/kernel/value": { + "dtype": "float32", + "sha256": "59a590bc9786236ff72db4845da3f8349aaaf073383a20a9c84b3e07526c04bc", + "shape": [ + 12, + 64, + 768 + ] + }, + "transformer/layers/9/self_attn/query/bias/value": { + "dtype": "float32", + "sha256": "ff0da73bb5d771687d86f20c9fb565ccc93d0101b8b59b5521fadc63f3d073d9", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query/kernel/value": { + "dtype": "float32", + "sha256": "05ecc00df4a2998ac2d85ceb9f0e6c3d6c425b6ea7f21aefa2620e75d896e476", + "shape": [ + 768, + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/query_ln/scale/value": { + "dtype": "float32", + "sha256": "ee62dead786a445f2bd87d41faba9132cd76962e93dc234fb10689958738f126", + "shape": [ + 64 + ] + }, + "transformer/layers/9/self_attn/value/bias/value": { + "dtype": "float32", + "sha256": "7087b4fe5ac8a46cfcab9e29c0c19624fe59172e9285b845b5fb15f534b9ca8f", + "shape": [ + 12, + 64 + ] + }, + "transformer/layers/9/self_attn/value/kernel/value": { + "dtype": "float32", + "sha256": "9974a5c45562310ea3b02117ad6cb5fd8153854d339e8b6b6b3b9003f2db5fe3", + "shape": [ + 768, + 12, + 64 + ] + }, + "vae/vae/decoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "1f485282296251d870b39b61d97e871e2d326f4e4ae2317dfa19c53910225bad", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "d42f4836883c41079ee6fa8b433418d6290e68bde820b1bad8ea1d07c4b8e611", + "shape": [ + 3, + 3, + 4, + 512 + ] + }, + "vae/vae/decoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "fd378bdd9d8c5ff306ea6ee50dd8a4dfa29f33ae7fd77516433083f4f7ecbc67", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "3e94ce9052653edb9ec5267ec8d4f233edf50b102a4b8a937a910a1f6d51528a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "924ed3c7be684daf687e7f28f56baced27bc2d9b69193bba4adb833d7efabfac", + "shape": [ + 3 + ] + }, + "vae/vae/decoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a59a09a27cdddc5c0fa72fb40f7272afeeb4512dfa021548617b493990563994", + "shape": [ + 3, + 3, + 128, + 3 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "c8463104a88fa0ae8842425b1e6c0aa6700dd315845ea677d5e27c5fcfea04d9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "4355032822e148fc76b5a72de0341d59682cd08638f1f742ecaaa0f94577b7b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "85a49f8bde4358125daf53ec14c2409d2c5bcdd16c38205cca3e9ad6b1a61346", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "040d54f3568cdfcf8274c4d3b7ba6b6668eba41e025b8ed5584b7ada475b625b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "b3a033c23ca569b85a17213d0f81dac5920eed618d2ffea58343117587f8faf7", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "a7cc59e3d58aff9ffd2f3776675e1bbfe5bbaeb217ea408acf04ea7806f5282b", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "0184aead285caab701ef2f0fa28a87345e53697ce481af8704744bf12de0b283", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "c749983771243f01b75913070efa45d6b84596a38c07dfb37ec8cf46e66ca822", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "851f0dc7e949218417c39f69cffad133dcdac1d7fd7da8eaa577f3d4f44a2ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "77453080b11b2edaa44121ceba9eac387e36fca0bbb28a3bda7ed51cf6129625", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "bc05fab622e6c7419b09b9e86d2441c7028c5b8e1402f173643989c1aea94639", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "b24968d292922139e8a2fc7484c907fa6b194754a9bdf40fd71fbb3fcbee68b6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "59b0162bcca002fca96f7196c2187149ea7416fed4345f43c88b38ea334c1e9f", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "4a4abf4ca5c85e0c1435c2f2775e39bc162cf65613a52cc1c6e2486330ea7c59", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "95aec6d80e7290ad10e02697976a0be44dec67ad4fb4f6760db34c46d6b50e90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f10da91594a42465807e044028a5050d3facab04c2ce473c825d0e2321b74295", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b2e21f036fa7c5a32994b712b5030839444d7f90261ffecd0e652a5db0b01f90", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ce16555de7a47ba62c9007f2eec765668239ccee6b2bedbdff3e667fabfde497", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "464fd559bdc14d38527894dceb7a1aaba42d3782ce325e14ac69b7a05490e066", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "666608a1ecafce314f9d968d9ccb503951038b867d8c34d39f0c47cfb9e5f2aa", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "adb49e4f44d607ff859c6e8f953f6606c1417bbc241150d1b3a78ce76ebd01fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "482ed775d9dd2505a7225154393109c20a1cd7aaabc6c5d64601629a83dfaa1b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "27553113d3337c0c675c68c8f793fd3c71fbd57fc9f686a951ab60a476547d5c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "608d04e1b6d5ce034792c44eb25b1dd17a1450bac3bef52c3dd8f8de03e228c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5872d9f16b7d0c416566509b4484239ab9cffc887776e1c035e06c01cdad0017", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "51a7a2162c5f44d6abe2e7223da773c929c50adc50ed5a7be9be804b6cec3bb0", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "235a739396d54514d59cbdeb6b74ea1ee56e1157f812dc9d4853b0977a570e70", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "66e90efa447b3ce33372e6d6b476253283cbefb8030ff55cc21110e73ee1bc8e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2d0ebb261baa8f3b8d28f8be2bbf90558027a44eb61c5604ceafa0ea78b15d86", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c6d00001d300712c40069de6add540d0e6c28d70a8fa90710a220ebf5d0f612d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5d65bc30ba4893171e8f0070adc504b22d6de9a5ecb7d0793b85ec912d73920d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dc71602f2e5aa2e48559cf15f374d084b36f4fbd4d5c2efa26892a2dc2f6934b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "f67427644ff8681e5b15a958795b63344cd035f5c710a82280243a2c8665da50", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "516549a836c3bca53fca49a11cba5f134d1a3c0a3a3a183ec9fe4607b37d520a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1e1cb132cd5dd05a7521113c43dee5f1902c30342aaba688403d6cbe6ca29e7c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "510bb20e3f07af15742bd26ac592c7207b37da96f758fbc0a671aa6ed2016f5f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "bf61f9dd7823c12462707ece3ba8c7988de31cc3ae25f8f5d1456debc8c1eb71", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f79a336f23d2494e42c4cc7dc2b0fa0d73a6681aaf8fc8994df94c51da2def2f", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f9287cc9c75d99954752ab39d8fac1d4cb24209cd938cd107ffaf8cc6bd27160", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "70bf2d4410856c02679ed4d9990fdc8be4a87bc8dca6c98e01ebf23ca2f1e86b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "dcddbbf2f50a4e6e3d467e34b0957f17706ed6852d63ada9f724026898e387b5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "9d1a8e347c4e33bc73107bf2094871a486831ee42d101f42ebba5e66bb791a3a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6cc582a85369f490c7519c7e5fde1c608c3014b0d3ca81bbb07cf8bdefff6721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cf8d97ecaf081233ecdd395c6573a3b62fd89d527b762c9e4d913cbc6c046381", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "0f3a408cda91e91edd4f187300eec219df015818c0e440df0efdff42d3c06d94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "852d24d773c3937318b309d00e877f57bee108d705b24b9f6304829be9194565", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6164a2612de78238a8d114ae8b921e8ee5f7d2aa3963ef9658b2f8ea54fd12e9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "a48683495e309df539c1c3e9696610ce47f86306811ea0dc44c0525a7eb54132", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "270f294a6a146afd396af5b274b68b72b9ecd3ad3d78466e3bccc03808c83028", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "25f35c3ef65afb973f5e434482ed217e0db0f73f8f1a31d37f6935b0784fa721", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "c95a90fd2cac1aaf6857f0dbfc4d76bd36a86d822003353a1357cf1e8847a1c6", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_0/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "e64f5d3718fbb8cdb3e3beda6e365748d64c27402757e102295e6c4f9846c80d", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "ff7a9812e07fc764ba3444d82ef23ec8c54b449eb6a1aa724314fd49d8320815", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5334aee9df72b791de9c3822dcbc6a432b7a7dabf4332da05e16ac517c293643", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2808c2e2a58eaa5397eb17d0b48d2096e16fbe019b742beeadb2a94d09dd4bbf", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6c941208ed362286de9eebdd19e07ae931e96d0c82b24590611855b85c8f1e43", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6c73223bbd956366dcd2e2d21e1ecfaeed7e2d251fa3bc2192fefd0a6819c697", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "10ae34f4005abfa31fdb53427c3f9f0df9b31b1a3f43cead335f87d420afd87c", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "faee4ff27bf8399afd5da1c77493697b1d8c51bd0795d87be68a36d51eb95c04", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "ddf0cb196577819ea6589b185abcf5ca099a4f233d08ff8ddbf0d67fdb681f53", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1121f2c2a7c05e09d71bd5b3b6f260bf833488c3a2da1a88b5efaf83ed307b2a", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "849f8ecbbb4f79ece6bfdd8a12b657ff4f02741ed453aac5857f1015752cd18b", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fd5a4ebdeccacd9dee369ee8f699fe6b60b88f1b2d9361bc7e66c540aac30d61", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "69cbdfeaec2688425952b89bd2424ad05a25c05abca7445bf5dc6c94f24c3f17", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "edd5d86d7f3df333496ea0559a255fc711c42da84688fc0fdbaed5e075dfcbe9", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b74095f523b1ba7c569923a836ee8281f1594e392952387e3aff5e39bc276762", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "e38bca8b93c2aeb66a41e9f4809640380454b55423231cfac042dfd2eb9e3ba5", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "f5f65409327f335eeb79d2f449264ecaaac0fff67ec9e8349be76043f251571b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "59ac3b6edfaf4e954b0e7fe6d29d5247036aa0244a19aefbad720c0cfa23de6b", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "c1b7150259e29d379c1017f0bc096c1bb970a1ff66983243ac5589d42284c848", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b8949fa468ea2695b12530e7a79c27cd6aa687d44c5cac45f0712b95b504d574", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "39b1068a3b8036240ea9057557b61c6aa8d6aa3ce80e0a3fde2cb550b2b2a2db", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "4acf532372dc21fc7958938f3de70ac54ecacdd601c29b76c543bfd48b7d3ef8", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "dbb8641fbf87865f1b5246f328f170f7c22e21c705f6950a8ad218bef91e3155", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "b019e13fa4a707ea61234541b6f8d7d255dc80c9ebeee1a57970c492a7d556c2", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "41048552fb395b8888cf89dd9a83e0a124b793cd955063c2a987e8bf22578502", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "39e7e8fd87edca23cba6ba099cadac54ba277729b72fe7dd31ab0da5a8fe28fd", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_1/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "2a8326c7cabfe79390e0155d682868d982ad8fdff390c50d5ebbf42fb5e0ca36", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "07cda010466e44977bce0765fe0a38e1a4b4472c2683f1207083ec00f1212ba8", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "ce818aa9e8a0c8be7e66ec8e9186f2bdfb04824eed17278d993a94c4fd37642d", + "shape": [ + 3, + 3, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "82b69d8e6a4c8479d9638885e5611de41513e5e8ac5b501438876a492712a8b9", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2352f238d878f44ed353d33c46c9d927f5ce67ec203ea3409a4f9f52f365fc22", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "53061bd8b75c78911529dc7c93824e0be30ef2adffc496ef93aea6de2db9bbfd", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "610ae868d246864bb5c88e5216be51fd2a03ac7529e8f1757d040b52cd67802c", + "shape": [ + 1, + 1, + 512, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "dabddacf88e8b00535552d37b7ecdbcbcdfab13a6afb457c83e6b222320e377d", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "2735929b7c4d1dd315248fa55ecb112471eabbd5995b663be331e81e114a9e94", + "shape": [ + 512 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "5b60ff520f2e4723278a62a1395e641911aa49c1dd9c70cb0a652426b7416ac0", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "680060e291a4178751a35d4776aa07c62fc6f8fb1807a635719de0a08a895e4b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "66ef8e7a42120dc1a3d6b4d7b442418a7c313001b7cbbdec0d77ebad19a992ac", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "8c36797641e2689b18c9a1e0f27cc31ed22f1270c82b1987e792332534a774d8", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "467f5f671801f27522ca1059d346ce81a4fcb6f499600932a90d7ea3547f2a93", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "b5b859023fbe4933890ddd35fa5305ff048f0ba0e821d0f60ef0b5e219f62518", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "6663bc4ca6a151927313e302db7a6f5c96c2fb82e34dcf8cbba513fb31d9e527", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "851e6e13039bd63e837e8eb7146ae697bc19f6be7a863101ae926192cef9f850", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "40ae9bcf58a7d01d60bac91e428df8e29fb962251c53f75d80491db71ba93a52", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b4ea5110211a473da86f2aa2dca7383fefbe4fea3a156529b4cf9a109872b0d7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2c45c9678c7cdd98170400b46e7264dff3ed9a984515ef11984798ca04f1c3bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "f4a87bd366aae198a6d480acc027c0cfdc5a1ee99ca1667fb9afab1b4c2e58f7", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "5868d5928588828cf4aa3592ffb90a06c7d654d72f7b9b52c69fd2cd66398bb7", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "88d27a08aaf128027eef67b57704454d7b3b9608089c2f54117e71a09d8db7cd", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "b1c03e766c6818e109f2e3b8bafa56f6cff665646f3dea0e6ae71c5e4a92a2b4", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "e04c0b78a73828318ba967b588f539a3acc0bf1d8e87f7eda6f64ae8a07b441a", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "317da4ae580133e3e0f544568ca4509149a07bda5b96d75346c3d30ec5015300", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "a33c429b13fc7b203b9e41c65fbd429f7925f3a40673c343724e9d0b7f50278f", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "86bca1f24c527eec5820cf9b7326fa8095fcba128ce715a47c6b06036579c0bf", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_2/upsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "61eb39f8df6ba73cf30e15cddf21a750e1af9ba3e08c8e71430dc060e8329a0b", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "1f969b858a9ec702ca53d237014a5506af34c0f912f7a3cf88cc8c7cba3cac82", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "5a7f41d1a89dc3e8f335ad92cc4fd18413bee9523c9c5379a9c7e60835b21c89", + "shape": [ + 3, + 3, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dfc42f1b2d11b3b12f360bf5c7fd3e41bf968dad573215b1106f8a35aca1c20d", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e80a63a3f326ca86dae89d16adf867ea2801b3d2e8ae0388eb641d3427a9c256", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "a1feb686734750b73d9ef0bdfde3dfe15d2e24ef5f07bea18547ee9326d6c234", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "4eb85305f736e2f4a2b6ab2bffc080a4dd9a0cdd6d7e5b43b01ee0a0b4447340", + "shape": [ + 1, + 1, + 256, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "3d26d070a07e732739a2cd158f4cb2729f47a21dcf23c6a18820954037f8328b", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "171c68b26753cf92771da6e25f607aa18df2bad417d10ae723c8e12abde9abf3", + "shape": [ + 256 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "73a1a80656595235daca0cc8e0cd53129cc44ec00bc8011282695355b3d7e386", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6d191496aa5cb64d8159cb21f265244f011a8abf74e782b9c997ecf8c753d316", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "6fecddd3973bf159b270c6f9030341df29a12edb5cb541862fc29d88b81be62b", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e2c941e23720a448e6d6e0a684965f1a229876e9e2e6d4ec03250b2c0a21cd", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "dbdea91e5a280572f2b97f2a6d2248a208adda2aec5fd90f923ce0fcc1a8962c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "c8c31d08590d05acde6be81ebe6a51ad3726d6d029beb9684dbff2e0167ceda4", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "5ae64884b642b9846c3b73605cdeadb373883dc55e2b57f353e591b95ebb2147", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "fcd6095f9babca03588569534a6b29d663d6b9eb8cfba9038bf0177f3930aee0", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aaa15421b93ed4333af1a5e9766b6775f4d46cf9a093e206ced9d615662ae86a", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5d39213b0174c8dcf2c7d2d4285d3b5d9ec3df5b454195fe308c2b04b33d3ef4", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "d3f0239f20c9836339c17c6d1a99ac1027a4f0f65de705b51c1ecfea65ef4a21", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "40599a423dfb5443ad1358922c8d3138cf65221c70ce507d8aeea04a96c5d8f6", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "304d7321ebfef6ab5fd6f395cb8ce60289c126cc5ca869c38893b232a2a8471c", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "6fd1c0a123a8e56dd2fff8748fafe5b08e985dd33f66c14b9bccc66536a88c2c", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "99641f31ee40b3eee17925f8c4529762fa52fc0f629bc344a4671796f1e2c11e", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "f32f7f17e99eaa796f92be818bd831d6274a266200940cb0e95d26bf60fab2d6", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "a94822386aea09678a734a1438f6ef944de9d88493b94dcbcc4330706d5fa910", + "shape": [ + 128 + ] + }, + "vae/vae/decoder/up_blocks_3/resnets_2/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "52d5122fbd4a5066ce62fefe35a93ca5f97e39ff9105713b28fef393586855f4", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/bias/value": { + "dtype": "bfloat16", + "sha256": "6331ecac159bb49673c473c6f337f251f8d9c31d205f02140c83466c92bb4239", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/conv_in/kernel/value": { + "dtype": "bfloat16", + "sha256": "01c0b06bec8905338cfd0c14944373d23bb5ed5f1efdf71ab54309d3aaf43b06", + "shape": [ + 3, + 3, + 3, + 128 + ] + }, + "vae/vae/encoder/conv_norm_out/bias/value": { + "dtype": "bfloat16", + "sha256": "174cee769dd4bebded94bb7d1cbc5aef93c6b02447136a02c93919538184a2ce", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_norm_out/scale/value": { + "dtype": "bfloat16", + "sha256": "e9c99b8ca8b76677ee092c4b8b7f520502c67b8da314fa796c41ef87f9578f79", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/conv_out/bias/value": { + "dtype": "bfloat16", + "sha256": "8a14d2241bb4b35b4505626e01463e91b8e17e93c91cc82349988bf43155ba98", + "shape": [ + 8 + ] + }, + "vae/vae/encoder/conv_out/kernel/value": { + "dtype": "bfloat16", + "sha256": "a88b18ca61cd00ba1225848d879a1a0b5fc3c27957ca414f8b0f353bd26fde7e", + "shape": [ + 3, + 3, + 512, + 8 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "773f2e3a6da207482df7eaa323a84b6a3b5b6c843f70ac0f6be903f4a0ae520e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "bdf36c07acddf9a2a4b5b3132f90f41bde81d7e921692acd2462fe232ac30541", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "99959f88d38f7835166b8f9e76444fc008d932b902a2c34b4a994f0341cc2493", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e0e81ee67d491992f729d795dd2e178ef13bbdf923fc8273508e5609362ceac3", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "b1ef48c3e6b2eb6cbedae470daf96aa18861dc416ee6f42a0b841bc4d2f84298", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "38c04f452a500d8c07f84bfadbf7e48e568c334d96e441071d351231655534af", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "23aff75e43e6d4b8e39ac6e7df1f8403fa36c1c309f9e0022c9ad6f8586adfab", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "ebf855617b538c1bec86a5b8741682b436fb32616088e6f7c52082cc14e75322", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "bb083c0aaa4a7a95e7353c47b622dd7150ed4073e8f4d8ff10057e15a4abf1c0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2dbaaa570ed5e36ba7c089d6e1edce070210050b6d9e5083f80879e371901b14", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "36c8d3054f504dac7ae718044c6b4ca999bfae582707f89d8920a0434b41b010", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "99a0aca2ca91217a9129ed8c53a24164cce7bf835969b7d905319c5944c88f98", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "31b6f15e1018570b232b63540dfb26891b04105405b3179e88c4e908214b389e", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bd2dd98720bf6e9ef15bb157a7dfc9787cbdefa18d7a961e23145923c94d3e0b", + "shape": [ + 3, + 3, + 128, + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7b0e142fd9c1b0905368df47ed43d9c1edc855d95080c68671471859a3046ec0", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "071c0ed6d4b684429337e1d5d934f94109a86700374a6e128ac41fa6404b89ae", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "44db5d34f6944e8bb7b122eb83c61f94e9b68c51d44871b2ea11c4f482c60460", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_0/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "6796305efb2d230bbe721c7c9af679a775e50814840ef3a787845f154197adea", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "10b27f99656ee3b07f077345eeb9d10f91782acca9a5137b4057b0133bf1fc26", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "92b1182ce8821ef31e3e233dd8c68213a6185a90a02d724c9f7759be923eb866", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "b7af570c368cd92d2635262ef3b777f515c29aa66002fcd069fa8e349097e6c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "e1cc57a207138702be84d6c2a0cb60523a091c69f094d0277d4dfb6035beb01f", + "shape": [ + 3, + 3, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "78408f60d8b2aff08d0f8acf8f5d316d97d804c2bc4f5033c7ad9fd7120120d1", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "665f87135983ecf2e2b12885a354becf245e7b9fc07fb1481f9e7f461fa6d61d", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "daabf25d95b84b711f4c531e65e17e76b38637498941baeeaea26f833101b7c4", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "77fded41464439100eebe3d089752e198877b2feb809636e79ff229dc10d0580", + "shape": [ + 1, + 1, + 128, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "f25fa7984bab80720fa7072d5c618bb1789cecb6274a5648abe22cd19cad40a1", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "11842f2ebcce2d5ec73eca311462917eb15171105fce6faec0714baa5de6e5c5", + "shape": [ + 128 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "fc632ead1adc9b0951b11bd90f10d6a28c892163e254f65250a8a008d5489249", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0ff76807ca7ac998478fcc09289ac43598f50b50545ea5d35a37472549f41f44", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "2e2ca33d68c93a579db99e9883c0130fdf7cbdb1f8bdbf6343800f3335072ab5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "cbc7aab9ef139073c8de640e1b3e6a0b80184427d71422f4a12b6da183beb9ac", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "6ce6133e34058e12ddbb7a2ab2f7ba69b8501460bfbbb1eb7875c1045fbf6b8d", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "e16946b17ce8c83f6e306dc87b331859a1388efa3e8baf4b533d69605c1bd936", + "shape": [ + 3, + 3, + 256, + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "1fb9d2d4e6bdf6bdd7840bd2487cdbb1874a1f4482cabd762f687f4d4c2f457f", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "25fba6f1dc356c6f4c36b01ed36c33626d4efe4cea9600ff53858ed0add7301a", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "aeb39dd95b769bdf7c76752581a2ae8faab3656acadcf9fad54c2e651d8932c5", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_1/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "2aa7e790089d3b9edd05cb5583ffa519a1e7df61f70302f89d82ddde49cde67c", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4b5005d43499814dce2777ae118846634e9881ee55b0ed57ee8486b28ec7f347", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/downsamplers_0/conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "3b9eeae2d4170a3f45a4372ddf6e5b57ab7efc3ca76ba4b8fb6ed87cf7e48cb1", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "3ed4121f48f80c3932512bb39685a8ca947d612545d113e28f9edda701c39c23", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "068f13e59b583aa618ef1e8019ad358547f0726644fb1cea8bccf9179a66fcda", + "shape": [ + 3, + 3, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "eef227da5c0de539318dbac03cec991c09bc18eeeac1feec14d587c9e70a8ad0", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "d8c7efda6c366e1a7aecd4211869f2c61dff073c15459e983f35110ee4fd0c0a", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/bias/value": { + "dtype": "bfloat16", + "sha256": "7d68c025c37f7f6fa795294e0da68c926faf17314947df2865aed43d0a8780e6", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/conv_shortcut/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd48e1e66eb4a8a765a7eef3c78493d2b5b8fd3c58e11ace9d3133bba1761602", + "shape": [ + 1, + 1, + 256, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "a45190df376d524d5dadb96501a528267525b7177b02c6edd2cd314e5ce4e948", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "b0e1df323fd990591a50fc7fb3a20875ded046c48dfd3125ccb663395503a8ac", + "shape": [ + 256 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1c2341de351577991631512d9b60bc32aecf0b87d5b9561f4e6adccc5b023bbe", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "b49e759d6601e33f7ee664708a83ee425bdc77ec276ee29e7bca27c85445113f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "647ca854797a36a0b770dcb35edf1e5db1b39960e0e5c630b1bdfcbb5e10439f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "6f02228dffdfe6bcb2441e8e52b261956a15cbe1dd7dbd1a540575cbbbdbab11", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "fcbe771bf543a980b71e8687d07d67ccdc868ce7e594fb0de1a5cd2479d49446", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "2e1595737c80c0c659151c2fc8e83fbebd52ec674f899f50782d95bc3ac88ba8", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7a13caa90cc79237a0fdbb883634f4cd9f67b2ea7542b10eec621727ab58fec9", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "5610b3e6e50098eaf3f7519730826351a3cd0369b1011461eface75a7d8910de", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "1646b4c52091713a1d308c83015a7f19b8417cbd4a89c568afd8a8a05e009cfb", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_2/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "d039692c7a3b833ab325a83469966e0a511e9f179165f56dcc8a35b542333145", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "65ab336fb59186fe8f70bcd96c1fc7361f06170c31f44dd57fc7eaf43334502f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "2d26181934bc89a10b658c1ad916156de831ae28ca78db0883b23d833f700d20", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "c62181f2875274ac5580137ff0da01a59d68cc3721485ea261356c4d6c762908", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "f5e7c85c7438291bf167928e1b9449bae4ffcfde4ecda24502dc660f76beea0e", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "26c748fb7ad7d696f4e0d68203a2ea3d8c1c671a69f9c302c575504a12ae365a", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "d7b4c89c04a1d35a1079365674f9b2ca3148ee207ebf69f5ab61034c81cf1449", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "20910fec7b7b2a3ae9b9d61a111698a8c0ba5982f8f2e624454efddfc2c54979", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "5c24f1783979be3b3b7a2b5c618ab31607baa52ee6aa7e20afb992fef7c0bd7b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "11aac69c34fcecb8bbc2ce0f964e3a93bec1d84566d3c021179427173c6d0481", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "0c5cc98cd8e7bcef205b0cb4552da4d3f3ce1a097a77926dd780d6fa600138f0", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "e2667f7b995f966aa2ec994ded2d3bfc4f025c7b3faca64064029edc6e6eea84", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "bfc8407ca8ae94f9d36605c8a3b73177c38176ba60ce1fd76483e4e67ce2faad", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "67f4e4bfd34de85b1b4493eb83d3eab2697f03c314bb50cb9a857b950c38659f", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "9298310029be409e1cd29afa89c65ab649d921669dd19505b64d23c9af262381", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "ac56041dddad73a8513549fc975f719afd66836b620ec14e12c3a727d17327b5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/down_blocks_3/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "df6a7c6f8a0529115919bf71b7ad29b8efbdb0db8903967e21a5638048bf7aba", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/bias/value": { + "dtype": "bfloat16", + "sha256": "346b46f8d13fe846faac607688b4775efcd271bafe7c2b27ed6d697d0ccffb11", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/group_norm/scale/value": { + "dtype": "bfloat16", + "sha256": "55da72fa6dd00bc01ef62896c8d8069f5e3cffff08f4b2b92c922a4d8f2a89fa", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/bias/value": { + "dtype": "bfloat16", + "sha256": "d2b445ea23cefe7e31e174e51524bb340c65511f8ce3e7940c2a79d9485e34ab", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/key/kernel/value": { + "dtype": "bfloat16", + "sha256": "1ab2988be6ac0725216e94d258a5f0289edd2249cce146ec21ef6cde2655ca85", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/bias/value": { + "dtype": "bfloat16", + "sha256": "585a7cd36de51c2c36b6774250e464c14c812d4e19447c295107cf20e300b10b", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/proj_attn/kernel/value": { + "dtype": "bfloat16", + "sha256": "6d54061a8695a6f7325716605834eb80b4f74d75128345edf2e35ac00ad25d7e", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/bias/value": { + "dtype": "bfloat16", + "sha256": "f5cda52e23dbe375f71f6305534b0c177c5fa8bc6f268ba96b5d6115f6c74f76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/query/kernel/value": { + "dtype": "bfloat16", + "sha256": "521db00246656c4768b1504d0a0d911992ebecf08e4016cada458398a7c84240", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/bias/value": { + "dtype": "bfloat16", + "sha256": "50889dfe8db29f82b6b372f626ac5fcd2ae762f55f055dd74ec1dbbfc45c9f37", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/attentions_0/value/kernel/value": { + "dtype": "bfloat16", + "sha256": "698879f6d6616ca4ce250a5e9bcedd7fbb59298c630dfaaf68e0df8ac601ae3f", + "shape": [ + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "a4d63b3e8e40520b5bade09abf1d0a3c7766cb873c48a02a964693905aa37429", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "12d20d065cbe97ad0fcc22c50c78526ba241e90bd47f595db79d48a50d62ad5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "2e8d3223514a535d59e4e9fbaeb2f9d4f94fd0dd2a63680affe946bac6ed2891", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "fd6ee2ae2fedd47b275173100a29e6ecfeaf5388c0576d79820c93297000c9c4", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "7f13a79e45aeb920cd4545170ea1315812355501976a1dc7c8cf69e5b39a53c4", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "4ea29842c2174fec60622ef3c7bfce852f5c1b299b25136926f7eaf76e052b73", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "c9775b0d46bbf241a378759b33cf561c11976ce4fd6c00daec8f4b8afe6c47a5", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_0/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "821ad2158cce7e4a2628ecd86fce9094df2a5e15c0fb81225c2c8b053fef4270", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/bias/value": { + "dtype": "bfloat16", + "sha256": "89ad6a93cbb5fdc3f4eb58b71cb9ffd1fce2830aebae7cbeab749fa019e53e76", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv1/kernel/value": { + "dtype": "bfloat16", + "sha256": "24e4d5e7e9d66be0cf4944107759159c029ac5017259ed37208201c11cd384f6", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/bias/value": { + "dtype": "bfloat16", + "sha256": "8e22c951ed5f549e6a4f67ca7c895a8afcf040345cbf2e1b4e7b9e6905ee40a8", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/conv2/kernel/value": { + "dtype": "bfloat16", + "sha256": "658719ac5a6c506d7af1497c8749f9831f01a16250bf101d3eaea253a659cf5c", + "shape": [ + 3, + 3, + 512, + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/bias/value": { + "dtype": "bfloat16", + "sha256": "04e74bc62b155f2b2b50043ffb3164f2ea080367497dce99c77e8d703ef65151", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm1/scale/value": { + "dtype": "bfloat16", + "sha256": "62b6918c190d4cbe121138903e1645cf0808df71b9f3ed2bb2112383ebf0a1d1", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/bias/value": { + "dtype": "bfloat16", + "sha256": "4a7ca183772d45bad0f01c410a47458d000fc202c504ca77a1928a049b41f401", + "shape": [ + 512 + ] + }, + "vae/vae/encoder/mid_block/resnets_1/norm2/scale/value": { + "dtype": "bfloat16", + "sha256": "0e3294d478784bee139778313a89cbfb2582616e11a94cb11474a991429f54d3", + "shape": [ + 512 + ] + }, + "vae/vae/post_quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "4dcc3ed96acb9312b54714af5fbb347ade056de4bc4f30f1a02cbba016eff96a", + "shape": [ + 4 + ] + }, + "vae/vae/post_quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "74010fb719a48ca0bd4b78660ce2249acb9f64b60f0fcdebf360848eaae9363e", + "shape": [ + 1, + 1, + 4, + 4 + ] + }, + "vae/vae/quant_conv/bias/value": { + "dtype": "bfloat16", + "sha256": "37cb4917563f4ee562fe05229d34a1bb8cf853a113e4f95eca53d70b9a6ede07", + "shape": [ + 8 + ] + }, + "vae/vae/quant_conv/kernel/value": { + "dtype": "bfloat16", + "sha256": "6b5d8e26e76f48ff5719edc8e0c3d766fdc28bf7f07230023346e7eb446ba8d8", + "shape": [ + 1, + 1, + 8, + 8 + ] + }, + "vae/vae/to_nnx__rngs/default/count/value": { + "dtype": "uint32", + "sha256": "d22025109bb7d889a919e8c3608b91821caaee3242ce2add8c5e2efc9b8a1b3c", + "shape": [] + }, + "vae/vae/to_nnx__rngs/default/key/value": { + "dtype": "key", + "sha256": "3b658cb4fc28afb4b5d743172e34dc20f3c052d7360c528f555ed8cd4a45f119", + "shape": [] + } + }, + "root_sha256": "932ab1e9013d756c3a8c80c1c7c3d9c6df2f6bc46426733805fabff9d1839460" + }, + "vae": { + "image_shape": [ + 256, + 256 + ], + "latent_shape": [ + 32, + 32, + 4 + ], + "scale_latents": true, + "source_sha256": "084f0c7cbb60b1f4d40bd328fd3c544fac7f46c7dafa0f146d68915bb577eb72", + "type": "StableVAE" + }, + "weights": { + "ema_decay": 0.995, + "ema_update_rate": 0.005, + "source": "ema", + "source_git_commit": "37f75b517bde668295c8f750b2d98e777bb35798", + "source_training_step": 40000, + "training_contract_sha256": "f48d843e060cba4ec72237e7a5fff5cc406cbdac4ea14e63421c28eac0824c70" + } +}